// Helper function for reduction ops.
//
// input_shape: 1-D Tensor, the shape of the Tensor being reduced.
// axes: 1-D Tensor, the reduction axes.
//   Note that the reduction indices are in the range
//   -rank(input_shape), rank(input_shape)
// returns a 1-D Tensor, the output shape as if keep_dims were set to True.
Output ReducedShapeHelper(const Scope& scope, const Output& input_shape,
                          const Output& reduction_axes);

}

void ConvertToMlirShape(const TensorShape& input_shape,
                        llvm::SmallVectorImpl<int64_t>* shape) {
  shape->reserve(input_shape.dims());
  for (const auto& d : input_shape) {
    shape->push_back(d.size == kTFDynamicSize ? ShapedType::kDynamic : d.size);
  }
}

Status ConvertToMlirShape(const TensorShapeProto& input_shape,
                          llvm::SmallVectorImpl<int64_t>* shape) {

    ArrayRef<int64_t> input_shape = input_type.getShape();
    if (reshape_shape.size() > input_shape.size()) return failure();

    // Extend the input shape with leading 1s to match the broadcast shape.
    ArrayRef<int64_t> broadcast_shape = output_type.getShape();
    SmallVector<int64_t, 4> input_shape_extended;
    input_shape_extended.append(broadcast_shape.size() - input_shape.size(), 1);

      "metadata": {
        "id": "rTcHwDPBPchd"
      },
      "outputs": [],
      "source": [
        "input_shape = (1, 224, 224, 3)\n",
        "\n",
        "jax_callable = jax2tf.convert(\n",
        "    ResNet50(\n",
        "      input_shape=input_shape[1:],\n",
        "      pooling='avg',\n",
        "  ).call,\n",
        "    with_gradient=False,\n",

// InferShapeForFunction.
FailureOr<bool> InferModuleShape(ModuleOp module, int64_t max_iterations = 10,
                                 ArrayRef<TypeID> ops_to_skip = {},
                                 ArrayRef<ArrayRef<int64_t>> input_shapes = {});

// Given a tensorflow NodeShape string, returns a vector of argument shapes
// that can be used with InferShapeForFunction.
// TF NodeShape uses `,` to separate dimensions, and `:` to separate arguments.

// Converts an TensorFlow shape to the one used in MLIR.
void ConvertToMlirShape(const TensorShape& input_shape,
                        llvm::SmallVectorImpl<int64_t>* shape);

// Converts an TensorFlow shape proto to the one used in MLIR.
Status ConvertToMlirShape(const TensorShapeProto& input_shape,
                          llvm::SmallVectorImpl<int64_t>* shape);

                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("b", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("c", "Placeholder", {},
                               {{"dtype", DT_FLOAT}, {"shape", input_shape}});
  *graph.add_node() = MakeNode("m", "TestFn", {"a", "b", "c"}, {});
  return graph;
}

  return func;
}

func::FuncOp createMaxUnpoolingFunc(
    mlir::Builder* builder, const SmallVector<int64_t, 4>& input_shape,
    const SmallVector<int64_t, 4>& output_shape) {
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());

// Please see the implementation file for documentation of these options.

// Import options.
extern llvm::cl::opt<std::string> input_arrays;
extern llvm::cl::opt<std::string> input_dtypes;
extern llvm::cl::opt<std::string> input_shapes;
extern llvm::cl::opt<std::string> output_arrays;
extern llvm::cl::opt<std::string> control_output_arrays;
extern llvm::cl::opt<std::string> inference_type;
extern llvm::cl::opt<std::string> min_values;

      enable_soft_placement,  set_original_tf_func_name};

  auto module_or = tensorflow::GraphdefToMlirTranslateFunction(
      input, input_arrays, input_dtypes, input_shapes, output_arrays,
      control_output_arrays, options, context);
  if (!module_or.status().ok()) return nullptr;
  return std::move(module_or).value();
}

static TranslateToMLIRRegistration GraphdefToMlirTranslate(