ArrayRef<int64_t> in_shape = ranked_type.getShape();
  if (in_shape.empty() || in_shape[0] < 0) {
    return context_op->emitOpError()
           << "A map_outside_compilation op's input and output shapes must "
              "have rank at least one and the first dimension must be known.";
  }
  int64_t split_size = in_shape[0] / num_cores_per_replica;
  if (in_shape[0] % num_cores_per_replica != 0) {

  auto x_shape = Shape(scope, x);
  auto output_shape = Shape(scope, op.output(0));

  // Reduce away broadcasted leading dims.
  auto reduce_x = internal::BroadcastGradientArgs(scope, x_shape, output_shape);
  auto gx_sum =
      ReduceSum(scope, gx, /*axis=*/reduce_x.r0, ReduceSum::KeepDims(true));
  auto gx_sum_reshape = Reshape(scope, gx_sum, x_shape);

  TF_SetStatus(status, TF_OK, "");
  auto* cc_ctx = reinterpret_cast<InferenceContext*>(ctx);
  auto* cc_result = reinterpret_cast<ShapeHandle*>(result);
  Status s = cc_ctx->Subshape(*reinterpret_cast<ShapeHandle*>(shape_handle),
                              start, end, cc_result);
  Set_TF_Status_from_Status(status, s);
}

int64_t TF_DimensionHandleValue(TF_DimensionHandle* dim_handle) {

  // https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-mat-mul
  int64_t x_row_dim = x_shape[x_shape.size() - 2];
  int64_t x_col_dim = x_shape[x_shape.size() - 1];
  int64_t y_row_dim = y_shape[y_shape.size() - 2];
  int64_t y_col_dim = y_shape[y_shape.size() - 1];
  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  size_t ByteSize() const override;
  void* Data() const override;
  bool IsAligned() const override;
  bool CanMove() const override;
  std::string SummarizeValue() const override;

  void SetShape(const int64_t* dims, int num_dims);
  Status ToTensor(tensorflow::Tensor* dst) const;
  Status BitcastFrom(const TensorInterface& from, DataType type,
                     const int64_t* new_dims, int num_new_dims);

  )
  def test_einsum_ptq_model(
      self,
      equation: str,
  ):
    _, y_shape, bias_shape, x_signature, y_signature = (
        self._prepare_sample_einsum_datashapes(equation, use_bias=True)
    )

    model = self._create_einsum_model(
        self._input_saved_model_path,
        equation,
        y_shape,
        x_signature,
        y_signature,
        bias_shape,
    )

  TF_DeleteExecutionContext(ctx);
  return wrap(func);
}

TF_AbstractTensor* TF_AddFunctionParameter(TF_ExecutionContext* func,
                                           TF_DataType dtype, TF_Shape shape,
                                           TF_Status* s) {
  DCHECK_GE(shape.num_dims, -1);
  TracingTensorHandle* t;
  TracingContext* tracing_ctx = dyn_cast<TracingContext>(unwrap(func));
  if (!tracing_ctx) {

  if (op.num_inputs() != 2) {
    return errors::InvalidArgument("BroadcastTo requires 2 inputs");
  }

  auto x_shape = Shape(scope, op.input(0));
  auto args = internal::BroadcastGradientArgs(scope, x_shape, op.input(1));
  auto sum_gx = Sum(scope, grad_inputs[0], args.r0);
  grad_outputs->push_back(Reshape(scope, sum_gx, x_shape));
  grad_outputs->push_back(NoGradient());
  return scope.status();
}

  auto shape = value_type.getShape();
  llvm::SmallVector<int64_t> sum_shape;
  for (int i = 0; i < shape.size(); ++i) {
    if (std::find(redux_axes.begin(), redux_axes.end(), i) ==
        redux_axes.end()) {
      sum_shape.push_back(shape[i]);
    }
  }
  return rewriter->create<TF::SumOp>(
      loc, RankedTensorType::get(sum_shape, value_type.getElementType()), value,
      redux_op);
}

        }

        shape_inference::ShapeHandle handle;
        TF_RETURN_IF_ERROR(
            context->MakeShapeFromPartialTensorShape(arg_shape.shape, &handle));
        TF_RETURN_IF_ERROR(shape_refiner->SetShape(n, 0, handle));
      }
    }

    // Sometimes we have VariableShape nodes in while loop (after Enter nodes).
    // They won't be constant-folded because TensorFlow constant folding does