element. For example, if you have a single image of shape `[height, width,
channels]`, you can make it a batch of 1 image with `expand_dims(image, 0)`,
which will make the shape `[1, height, width, channels]`.

Other examples:

```
# 't' is a tensor of shape [2]
shape(expand_dims(t, 0)) ==> [1, 2]
shape(expand_dims(t, 1)) ==> [2, 1]
shape(expand_dims(t, -1)) ==> [2, 1]

  %0 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<10x20>} : (tensor<10x20xf32>) -> tensor<10x20xf32>
  %1 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<?x20>} : (tensor<10x20xf32>) -> tensor<10x20xf32>
  // Failing case which should not be folded.
  // CHECK: %[[NF:.*]] = "tf.EnsureShape"(%arg0) <{shape = #tf_type.shape<20x10>}>
  %2 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<20x10>} : (tensor<10x20xf32>) -> tensor<20x10xf32>

    new_filter_shape.push_back(filter_shape[transpose_dims[i]]);
  }

  auto get_array_idx = [](ArrayRef<int64_t> shape, const int i, const int j,
                          const int k, const int l) -> int64_t {
    return (i * shape[1] * shape[2] * shape[3]) + (j * shape[2] * shape[3]) +
           (k * shape[3]) + l;
  };

  // Transpose the filter value.
  // TODO: b/336203735 - Use `DenseElementsTransposer` instead of manual

// because starting from version 7, in presence of shape polymorphism JAX will
// emit stablehlo.custom_call @shape_assertion to verify at compile time that
// the code is used with compatible actual shapes.
// TFLite runtime kernels support shape checking and shape inference to some
// extent, it is okay to remove the shape assertion in most scenarios. However

  NodeBuilder node_builder(absl::StrCat(call_node, "_key_placeholder"),
                           "Placeholder", opts.op_registry());
  TensorShapeProto shape;
  shape.add_dim()->set_size(2);
  return opts.WithAttr("shape", shape)
      .WithAttr("dtype", DT_STRING)
      .WithAttr("_host_compute_call_node", call_node)
      .FinalizeBuilder(&node_builder);
}

             .ok()) {
      return std::nullopt;
    }

    const PartialTensorShape shape = shapes[e->src_output()];
    if (!shape.IsFullyDefined()) {
      return std::nullopt;
    }

    results[e->dst_input()] = shape;
  }
  return results;
}

string host_compute_node_name(const string& original_oc_name) {

    // CHECK-NOT:  "tf.Shape"(%[[ARG_1]])
    // CHECK-NOT:  "tf.Shape"(%[[ARG_3]])
    // CHECK:      %[[ARG_0_SHAPE:[0-9]*]] = "tf.Shape"(%[[ARG_0]])
    // CHECK:      %[[ARG_2_SHAPE:[0-9]*]] = "tf.Shape"(%[[ARG_2]])

      GetTFLiteType(tensor_type.getElementType()).value();
  std::vector<int32_t> shape;
  if (tensor_type.hasRank()) {
    llvm::ArrayRef<int64_t> shape_ref = tensor_type.getShape();
    shape = std::vector<int32_t>(shape_ref.begin(), shape_ref.end());
  }

  variant_params.push_back(
      tflite::CreateVariantSubType(builder_, builder_.CreateVector(shape),
                                   tflite_element_type, tensor_type.hasRank()));

  // CHECK-LABEL:mul_with_high_dims_dynamic_shape_both_sides
  // CHECK: %[[SHAPE:.*]] = "tfl.shape"(%arg0) : (tensor<8x7x6x5x?x3x2x1xf32>) -> tensor<8xi64>
  // CHECK: %[[SHAPE_1:.*]] = "tfl.shape"(%arg1) : (tensor<?x3x2x1xf32>) -> tensor<4xi64>
  // CHECK: %[[BROADCAST_ARGS:.*]] = "tfl.broadcast_args"(%[[SHAPE]], %[[SHAPE_1]]) : (tensor<8xi64>, tensor<4xi64>) -> tensor<8xi64>

  %0 = "tfl.shape"(%arg1) : (tensor<8x7x6x5x?x3x2x1xf32>) -> tensor<8xi64>
  %1 = "tfl.shape"(%arg2) : (tensor<?x3x2x1xf32>) -> tensor<4xi64>
  %2 = "tfl.broadcast_args"(%0, %1) : (tensor<8xi64>, tensor<4xi64>) -> tensor<8xi64>
  %3 = "tfl.shape"(%arg0) : (tensor<8x7x6x5x?x3x2x1xi1>) -> tensor<8xi64>