code: def test_tensors_inferred_from_map(ray_start_regular_shared):
      # Test map.
      ds = ray.data.range(10, parallelism=10).map(
          lambda _: {"data": np.ones((4, 4))}
      )
      ds = ds.materialize()
      assert str(ds) == (
          "MaterializedDataset(\n"
          "   num_blocks=10,\n"
          "   num_rows=10,\n"
          "   schema={data: numpy.ndarray(shape=(4, 4), dtype=double)}\n"
          ")"
      )
  
      # Test map_batches.
      ds = ray.data.range(16, parallelism=4).map_batches(
          lambda _: {"data": np.ones((3, 4, 4))}, batch_size=2
      )
      ds = ds.materialize()
      assert str(ds) == (
          "MaterializedDataset(\n"
          "   num_blocks=4,\n"
          "   num_rows=24,\n"
          "   schema={data: numpy.ndarray(shape=(4, 4), dtype=double)}\n"
          ")"
      )
  
      # Test flat_map.
      ds = ray.data.range(10, parallelism=10).flat_map(
          lambda _: [
              {"data": np.ones((4, 4))},
              {"data": np.ones((4, 4))},
          ]
      )
      ds = ds.materialize()
      assert str(ds) == (
          "MaterializedDataset(\n"
          "   num_blocks=10,\n"
          "   num_rows=20,\n"
          "   schema={data: numpy.ndarray(shape=(4, 4), dtype=double)}\n"
          ")"
      )
  
      # Test map_batches ndarray column.
      ds = ray.data.range(16, parallelism=4).map_batches(
          lambda _: pd.DataFrame({"a": [np.ones((4, 4))] * 3}),
          batch_size=2,
      )
      ds = ds.materialize()
      assert str(ds) == (
          "MaterializedDataset(\n"
          "   num_blocks=4,\n"
          "   num_rows=24,\n"
          "   schema={a: numpy.ndarray(shape=(4, 4), dtype=float64)}\n"
          ")"
      )
  
      ds = ray.data.range(16, parallelism=4).map_batches(
          lambda _: pd.DataFrame(
              {"a": [np.ones((2, 2)), np.ones((3, 3))]}
          ),
          batch_size=2,
      )
      ds = ds.materialize()
      assert str(ds) == (
          "MaterializedDataset(\n"
          "   num_blocks=4,\n"
          "   num_rows=16,\n"
          "   schema={a: numpy.ndarray(shape=(None, None), dtype=float64)}\n"
          ")"
      )
  
imports: import time
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.air.util.tensor_extensions.utils import _create_possibly_ragged_ndarray
  from ray.data.block import BlockAccessor
  from ray.data.extensions.tensor_extension import ArrowTensorArray, ArrowTensorType, ArrowVariableShapedTensorArray, ArrowVariableShapedTensorType, TensorArray, TensorDtype
  from ray.data.tests.conftest import *
  from ray.data.tests.util import extract_values
  from ray.tests.conftest import *
label_desc: Use image node_id 6 for calling a global function with example usage: test_tensors_inferred_from_map(ray_start_regular_shared) without return types