code: def _ensure_matrix(volatility, dim, dtype):
      """Converts a volatility tensor to the right shape."""
      # Works only for static rank.
      rank = len(volatility.shape)
      if not rank:
          return tf.eye(dim, dtype=dtype) * volatility
      if rank == 1:
          return tf.linalg.tensor_diag(volatility)
      # It is of rank 2 at least
      return volatility
  
file_path: tf-quant-finance/tf_quant_finance/models/legacy/brownian_motion_utils.py
imports: import tensorflow.compat.v2
  from tf_quant_finance.math import gradient
label_desc: Use image node_id 9 for calling a global function with example usage: _ensure_matrix(volatility, dim, dtype) and returns: volatility, unknown, tf