code: def sequence_mask(data, valid_length, mask_value=0, axis=0):
      """Sets all elements outside the expected length of the sequence to a constant value.
  
      This function takes an n-dimensional input array of the form [MAX_LENGTH, batch_size, ...] or
      [batch_size, MAX_LENGTH, ...] and returns an array of the same shape.
  
      `axis` means the axis of the length dimension and can only be 0 or 1. If `axis` is 0,
      the data must have shape [MAX_LENGTH, batch_size, ...]. Otherwise (axis=1), the data must have
      shape [batch_size, MAX_LENGTH, ...].
  
      `valid_length` gives the length of each sequence. `valid_length` should be
      a 1D int array with positive ints and has dimension [batch_size,].
  
      Parameters
      ----------
      data : tvm.te.Tensor
          N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...]
          depending on the value of `axis`.
  
      valid_length : tvm.te.Tensor
          1-D with shape [batch_size,]
  
      mask_value : float, optional
          The masking value, default 0
  
      axis : int, optional
          axis of the length dimension, must be 0 or 1, default 0
  
      Returns
      -------
      output : tvm.te.Tensor
          N-D with shape [MAX_LENGTH, batch_size, ...] or [batch_size, MAX_LENGTH, ...]
          depending on the value of `axis`.
      """
  
      assert (
          len(data.shape) >= 2
      ), f"only support data.ndim >= 2, received data.shape = {data.shape}"
      assert axis in (
          0,
          1,
      ), f"only support axis = 0, 1, received axis = {axis}"
      return cpp.sequence_mask(data, valid_length, mask_value, axis)
  
imports: from __future__ import absolute_import
  import tvm
  from tvm import te, topi
  from tvm.te import hybrid
  from .None import cpp, tag
  from .utils import const_vector, make_idx, within_index
label_desc: Use image node_id 25 for calling a global function with example usage: sequence_mask(data, valid_length, mask_value, axis) and returns: cpp