code: def make_polynomial_basis(
      degree: types.IntTensor,
  ) -> Callable[..., types.RealTensor]:
      """Produces a callable from samples to polynomial basis for use in regression.
  
      The output callable accepts a scalar `Tensor` `t` and a `Tensor` `X` of
      shape `[num_samples, dim]`, computes a centered value
      `Y = X - mean(X, axis=0)` and outputs a `Tensor` of shape
      `[degree * dim, num_samples]`, where
      ```
      Z[i*j, k] = X[k, j]**(degree - i) * X[k, j]**i, 0<=i<degree - 1, 0<=j<dim
      ```
      For example, if `degree` and `dim` are both equal to 2, the polynomial basis
      is `1, X, X**2, Y, Y**2, X * Y, X**2 * Y, X * Y**2`, where `X` and `Y` are
      the spatial axes.
  
      #### Example
      ```python
      basis = tff.experimental.lsm_algorithm.make_polynomial_basis_v2(2)
      x = [[1.0], [2.0], [3.0], [4.0]]
      x = tf.expand_dims(x, axis=-1)
      basis(x, tf.constant(0, dtype=np.int32))
      # Expected result:
      [[ 1.  ,  1.  ,  1.  ,  1.  ], [-1.5 , -0.5 ,  0.5 ,  1.5 ],
      [ 2.25,  0.25,  0.25,  2.25]]
      ```
  
      Args:
        degree: An `int32` scalar `Tensor`. The degree of the desired polynomial
          basis.
  
      Returns:
        A callable from `Tensor`s of shape `[batch_size, num_samples, dim]` to
        `Tensor`s of shape `[batch_size, (degree + 1)**dim, num_samples]`.
      """
  
      def basis(
          sample_paths: types.RealTensor, time_index: types.IntTensor
      ) -> types.RealTensor:
          """Computes polynomial basis expansion at the given sample points.
  
          Args:
            sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
              either shape `[num_samples, num_times, dim]` or
              `[batch_size, num_samples, num_times, dim]`.
            time_index: An integer scalar `Tensor` that corresponds to the time
              coordinate at which the basis function is computed.
  
          Returns:
            A `Tensor`s of shape `[batch_size, (degree + 1)**dim, num_samples]`.
          """
          sample_paths = tf.convert_to_tensor(
              sample_paths, name="sample_paths"
          )
          if sample_paths.shape.rank == 3:
              sample_paths = tf.expand_dims(sample_paths, axis=0)
          shape = tf.shape(sample_paths)
          num_samples = shape[1]
          batch_size = shape[0]
          dim = sample_paths.shape[
              -1
          ]  # Dimension should statically known
          # Shape [batch_size, num_samples, 1, dim]
          slice_samples = tf.slice(
              sample_paths,
              [0, 0, time_index, 0],
              [batch_size, num_samples, 1, dim],
          )
          # Shape [batch_size, num_samples, 1, dim]
          samples_centered = slice_samples - tf.math.reduce_mean(
              slice_samples, axis=1, keepdims=True
          )
          grid = tf.range(degree + 1, dtype=samples_centered.dtype)
          # Creates a grid of 'power' expansions, i.e., a `Tensor` of shape
          # [(degree + 1)**dim, dim] with entries [k_1, .., k_dim] where
          ## 0 <= k_i <= dim.
          grid = tf.meshgrid(*(dim * [grid]))
          # Shape [(degree + 1)**3, dim]
          grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
          # `samples_centered` has shape [batch_size, num_samples, 1, dim],
          # `samples_centered**grid` has shape
          # `[batch_size, num_samples, (degree + 1)**dim, dim]`
          # so that the output shape is `[batch_size, num_samples, (degree + 1)**dim]`
          basis_expansion = tf.reduce_prod(
              samples_centered**grid, axis=-1
          )
          return tf.transpose(basis_expansion, [0, 2, 1])
  
      return basis
  
file_path: tf-quant-finance/tf_quant_finance/models/longstaff_schwartz/lsm.py
imports: from typing import Callable
  import tensorflow.compat.v2
  from tf_quant_finance import types
  from tf_quant_finance import utils
label_desc: Use image node_id 1 for calling a global function with example usage: make_polynomial_basis(degree) and returns: basis, tf