code: def testShapeInformationPeserved(self):
      @def_function.function
      def linspace(start, stop, num, axis):
          return math_ops.linspace_nd(start, stop, num=num, axis=axis)
  
      # Constant num and axis leads to preserved known shape.
      output_shape = linspace.get_concrete_function(
          start=tensor.TensorSpec(
              shape=[64, None], dtype=dtypes.float32
          ),
          stop=tensor.TensorSpec(
              shape=[64, None], dtype=dtypes.float32
          ),
          num=10,
          axis=-1,
      ).output_shapes
      expected_shape = (64, None, 10)
      self.assertEqual(output_shape, expected_shape)
  
file_path: tensorflow-upstream/tensorflow/python/ops/math_ops_linspace_test.py
imports: from distutils.version import LooseVersion
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import googletest
label_desc: Use image node_id 2 for calling the LinspaceTest obj's underlying member method code with example usage: obj.testShapeInformationPeserved() and returns: math_ops