returnOkStatus();

});

REGISTER_OP("TPUPartitionedInputV2")

.Input("inputs: N * T")

.Output("output: T")

.Attr("N: int >= 1")

.Attr("T: type")

.Attr("partition_dims: list(int)")

.Attr("is_packed: bool = false")

.SetShapeFn([](InferenceContext* c) {

DataType dtype;

TF_RETURN_IF_ERROR(c->GetAttr("T", &dtype));

std::vector<int> partition_dims;

TF_RETURN_IF_ERROR(c->GetAttr("partition_dims", &partition_dims));

bool is_packed;

TF_RETURN_IF_ERROR(c->GetAttr("is_packed", &is_packed));

int num_partitions =1;

for (constint& partition_dim : partition_dims) {

num_partitions *= partition_dim;

}

bool replicated = partition_dims.empty();

int num_inputs_expected = is_packed ?1 : num_partitions;

if (!((replicated && !is_packed) ||

(c->num_inputs() == num_inputs_expected))) {

returnerrors::InvalidArgument("Expected", num_inputs_expected,

" inputs, got", c->num_inputs(),".");

}elseif (c->num_inputs() ==0) {

returnerrors::InvalidArgument(

"Expected at least one input to TPUPartitionedInputV2.");

}

ShapeHandle output_shape;

if (dtype == DT_RESOURCE) {

ShapeHandle previous_shape_handle;

const std::vector<shape_inference::ShapeAndType>* shapes_and_types =

nullptr;

for (int i = c->num_inputs() -1; i >=0; --i) {

shapes_and_types = c->input_handle_shapes_and_types(i);

if (shapes_and_types) {

ShapeHandle shape_handle = shapes_and_types->at(0).shape;

if (!c->FullyDefined(shape_handle)) {

returnerrors::InvalidArgument("Inputs must have static shape,",

"input[", i,

"] has unknown dimension.");

}

if (i != c->num_inputs() -1) {

ShapeHandle tmp;

if (!c->Merge(shape_handle, previous_shape_handle, &tmp).ok()) {

returnerrors::InvalidArgument(

"Inputs must have the same shape.");

}

}else {

previous_shape_handle = shape_handle;

}

}

}

if (shapes_and_types) {

TF_ASSIGN_OR_RETURN(

output_shape,

_ComputeOutputShape(c, previous_shape_handle, partition_dims));

std::vector<shape_inference::ShapeAndType> output_shapes_and_types;

output_shapes_and_types.push_back(shape_inference::ShapeAndType(

output_shape, shapes_and_types->at(0).dtype));

c->set_output_handle_shapes_and_types(0, output_shapes_and_types);

}

}

if (!c->FullyDefined(output_shape)) {

TF_ASSIGN_OR_RETURN(

output_shape,_ComputeOutputShape(c, c->input(0), partition_dims));

}

c->set_output(0, output_shape);

returnOkStatus();

});

}// namespace tensorflow

returnOkStatus();

});

REGISTER_OP("TPUPartitionedOutputV2")

.Input("inputs: T")

.Output("output: num_splits * T")

.Attr("T: type")

.Attr("num_splits: int >= 1")

.Attr("partition_dims: list(int)")

.SetShapeFn([](InferenceContext* c) {

DataType dtype;

TF_RETURN_IF_ERROR(c->GetAttr("T", &dtype));

std::vector<int> partition_dims;

TF_RETURN_IF_ERROR(c->GetAttr("partition_dims", &partition_dims));

int num_splits;

TF_RETURN_IF_ERROR(c->GetAttr("num_splits", &num_splits));

if (dtype == DT_RESOURCE) {

returnerrors::Unimplemented("Not implemented.");

}elseif (c->num_inputs() ==0) {

returnerrors::InvalidArgument(

"Expected at least one input to TPUPartitionedOutputV2.");

}

ShapeHandle handle = c->input(0);

int rank =InferenceContext::Rank(handle);

int num_cores_per_replica =1;

for (constint& partition_dim : partition_dims) {

num_cores_per_replica *= partition_dim;

}

if (num_splits != num_cores_per_replica) {

returnerrors::InvalidArgument("Expected", num_cores_per_replica,

" splits.");

}elseif (rank > (int)partition_dims.size()) {

returnerrors::InvalidArgument("Expected at least", rank,

" partition dimensions.");

}

for (int i =0; i < rank; ++i) {

shape_inference::DimensionHandle dim;

TF_RETURN_WITH_CONTEXT_IF_ERROR(

c->Divide(c->Dim(handle, i), partition_dims[i],

true/* evenly_divisible*/, &dim),

"Number of ways to split should evenly divide the split dimension");

TF_CHECK_OK(c->ReplaceDim(handle, i, dim, &handle));

}

for (int i = num_splits -1; i >=0; --i) {

c->set_output(i, handle);

}

returnOkStatus();

});

}// namespace tensorflow