content: import logging
  from typing import List, Optional, Sequence, Union
  
  import numpy as np
  import tensorrt as trt
  import torch
  from tensorrt import ITensor as TRTTensor
  from torch.fx.node import Target
  from torch_tensorrt.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion import impl
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      broadcastable,
      cast_trt_tensor,
      get_positive_dim,
      get_trt_tensor,
      set_layer_name,
      to_numpy,
  )
  from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise
  from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
  from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  def select(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: int,
      index: int,
  ) -> TRTTensor:
      if not isinstance(input, TRTTensor):
          raise RuntimeError(
              f"slice_tensor received input {input} that is not part "
              "of the TensorRT region!"
          )
  
      ranks = len(input.shape)
      dim = get_positive_dim(dim, ranks)
  
      indices_tensor = get_trt_tensor(
          ctx, np.array(index, dtype=np.int32), f"{name}_indices_tensor"
      )
      layer = ctx.net.add_gather(input, indices_tensor, dim)
  
      return layer.get_output(0)
  
  
  def is_boolean_tensor(
      tensor: Union[TRTTensor, np.ndarray, torch.Tensor, torch.fx.Node],
  ) -> bool:
      if isinstance(tensor, torch.Tensor):
          return bool(tensor.dtype == torch.bool)
      elif isinstance(tensor, np.ndarray):
          return bool(tensor.dtype == np.bool_)
      elif isinstance(tensor, TRTTensor):
          return bool(tensor.dtype == trt.DataType.BOOL)
      # when index is a node
      else:
          val = tensor.meta.get("val")
          if val is not None and val.dtype is torch.bool:
              return True
  
      return False
  
  
  def expand_boolean_indices(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      indices: Sequence[Union[TRTTensor, np.ndarray, torch.Tensor]],
  ) -> Sequence[Union[TRTTensor, np.ndarray, torch.Tensor]]:
      new_indices = []
      for i, ind in enumerate(indices):
          if ind is not None and is_boolean_tensor(ind):
              _LOGGER.debug(
                  f"Boolean index detected at position {i}, converting with nonzero()"
              )
              mask_tensor = get_trt_tensor(ctx, ind, name + f"_bool_mask_{i}")
  
              nonzero_layer = ctx.net.add_non_zero(mask_tensor)
              set_layer_name(
                  nonzero_layer, target, name + f"_bool_nonzero_{i}", source_ir
              )
              nonzero_indices = nonzero_layer.get_output(0)
  
              # nonzero returns shape [N, dims], we need to extract dim i
              if len(indices) == 1:
                  # x[mask] — 1D mask
                  to_squeeze = nonzero_indices
              else:
                  # Advanced multi-axis mask: extract index i from shape [N, D]
                  gather_axis = 1  # dim index
                  gather_layer = ctx.net.add_gather(
                      nonzero_indices,
                      get_trt_tensor(ctx, i, name + f"_dim_index_{i}"),
                      gather_axis,
                  )
                  set_layer_name(
                      gather_layer, target, name + f"_bool_nonzero_extract_{i}", source_ir
                  )
                  to_squeeze = gather_layer.get_output(0)
              squeeze_layer = ctx.net.add_shuffle(to_squeeze)
              squeeze_layer.reshape_dims = (-1,)
              set_layer_name(
                  squeeze_layer,
                  target,
                  name + f"_bool_mask_squeeze_{i}",
                  source_ir,
              )
              squeezed_index = squeeze_layer.get_output(0)
              new_indices.append(squeezed_index)
          else:
              new_indices.append(ind)
      return new_indices
  
  
  def index(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      indices: Sequence[Union[TRTTensor, np.ndarray, torch.Tensor]],
  ) -> TRTTensor:
      adv_indx_indices = []
      tensor_indices = []
      # is_numpy is a flag to specify if all the indices are numpy or torchTensor.
      # If any is not this flag will be set to False
      _LOGGER.debug(
          "Determining whether aten.index constant-index optimization can be invoked"
      )
      indices = expand_boolean_indices(ctx, target, source_ir, name, input, indices)
      is_numpy = all(
          isinstance(ind, (torch.Tensor, np.ndarray))
          for ind in indices
          if ind is not None
      )
      # here we need to check if all the index are broadcastable
      # if no, then we need to broadcast
      last_index = None
      for i, ind in enumerate(indices):
          if ind is not None:
              _LOGGER.debug(f"Shape of {i} index is {ind.shape}")
              adv_indx_indices.append(i)
              # torch.nn.parameter.Parameter=> numpy array
              # numpy array is kept as numpy
              # other cases are kept as TRTTensor
              if is_numpy:
                  ind = to_numpy(ind)
              else:
                  ind = get_trt_tensor(ctx, ind, name + f"_parameter_to_fp32_tensor_{i}")
              if last_index is not None:
                  assert broadcastable(
                      ind, last_index
                  ), "The indices should be broadcastable!"
              last_index = ind
              tensor_indices.append(ind)
  
      if not tensor_indices:
          cast_layer = ctx.net.add_cast(input, trt.int32)
          set_layer_name(cast_layer, target, name + "_index_casted", source_ir)
          return cast_layer.get_output(0)
      elif len(tensor_indices) == 1:
          indices_tensor = get_trt_tensor(
              ctx, tensor_indices[0], name + "_parameter_to_fp32_tensor"
          )
          index = adv_indx_indices[0]
          _LOGGER.debug(f"The advanced index indices is {adv_indx_indices}")
          gather_layer = ctx.net.add_gather(input, indices_tensor, index)
          set_layer_name(gather_layer, target, name + "_index_gather", source_ir)
          return gather_layer.get_output(0)
      else:
          input_shape = input.shape
          _LOGGER.debug(f"The input shape is {input.shape}")
          rank = len(input_shape)
          adv_indx_count = len(adv_indx_indices)
          dim_tensor_list = []
  
          for i in range(rank):
              if input_shape[i] != DYNAMIC_DIM:
                  dim = input_shape[i]
                  dim_tensor = get_trt_tensor(ctx, dim, name + f"_individual_dim_{i}")
              else:
                  dim_tensor = get_shape(
                      ctx, target, source_ir, name + f"_individual_dim_dyn_{i}", input, i
                  )
              # dim_tensor_list is a list of tensors
              dim_tensor_list.append(dim_tensor)
  
          # for cases like
          # t: [x_1, y_1, y_2, ..., x_m, ..., y_n] -> t: [x_1, x_2, ..., x_m, y_1, y_2, ..., y_n],
          # where t is a tensor of rank m+n, {x_i} are axes where tensor index is provided, and {y_i} are axes
          # for ":"
          # Examples: x.shape = (10,20,30,40,50)
          # ind_1, ind_2 broadcasted to (2,3,4)
          # x[:, ind_1, ind_2] = 10, 2, 3, 4, 40, 50
          # x[:,ind_1, :, ind_2] = 2, 3, 4, 10, 30, 50
          transpose_layer = ctx.net.add_shuffle(input)
          new_order = []
          for i in range(adv_indx_count):
              new_order.append(adv_indx_indices[i])
          for i in range(rank):
              if i not in adv_indx_indices:
                  new_order.append(i)
          _LOGGER.debug(f"The new transpose order is {new_order}")
  
          transpose_layer.second_transpose = tuple(new_order)
          set_layer_name(transpose_layer, target, name + "_index_transpose", source_ir)
          transpose_tensor = transpose_layer.get_output(0)
  
          # Flatten [x_1, x_2,.......x_m, y_1, y_2,.....y_n]
          # transpose_tensor_shape = ctx.net.add_shape(transpose_tensor)
          transpose_tensor_shape = transpose_tensor.shape
          _LOGGER.debug(f"The shape of transpose tensor is {transpose_tensor_shape}")
  
          mult_d0 = 1
          dim_tensor_shape_mult_d0 = 1
          for i in range(adv_indx_count):
              if transpose_tensor_shape[i] == DYNAMIC_DIM:
                  dim_tensor_shape_mult_d0 = get_shape(
                      ctx,
                      target,
                      source_ir,
                      name + f"_transpose_tensor_shape_mult_d0_{i}",
                      transpose_tensor,
                      i,
                  )
              else:
                  dim_tensor_shape_mult_d0 = transpose_tensor_shape[i]
              mult_d0 = convert_binary_elementwise(
                  ctx,
                  target,
                  source_ir,
                  name + f"_shape_{i}",
                  trt.ElementWiseOperation.PROD,
                  mult_d0,
                  dim_tensor_shape_mult_d0,
              )
          mult_d1 = 1
          dim_tensor_shape_mult_d1 = 1
          for i in range(adv_indx_count, rank):
              if transpose_tensor_shape[i] == DYNAMIC_DIM:
                  dim_tensor_shape_mult_d1 = get_shape(
                      ctx,
                      target,
                      source_ir,
                      name + f"_transpose_tensor_shape_mult_d0_{i}",
                      transpose_tensor,
                      i,
                  )
              else:
                  dim_tensor_shape_mult_d1 = transpose_tensor_shape[i]
  
              mult_d1 = convert_binary_elementwise(
                  ctx,
                  target,
                  source_ir,
                  name + f"_shape_{i}",
                  trt.ElementWiseOperation.PROD,
                  mult_d1,
                  dim_tensor_shape_mult_d1,
              )
  
          concat_tensor_layer = ctx.net.add_concatenation(
              [
                  get_trt_tensor(ctx, mult_d0, name + "_d0_shape"),
                  get_trt_tensor(ctx, mult_d1, name + "_d1_shape"),
              ]
          )
          set_layer_name(concat_tensor_layer, target, name + "_index_Concat", source_ir)
          concat_tensor = concat_tensor_layer.get_output(0)
  
          reshape_layer = ctx.net.add_shuffle(transpose_tensor)
          reshape_layer.set_input(1, concat_tensor)
          flatten_tensor = reshape_layer.get_output(0)
  
          _LOGGER.debug(f"The flatten tensor shape is {flatten_tensor.shape}")
  
          # tensor index = \sum_{i=1}^m (ind_i * \prod_{j=i+1}^m (x_j)),  ind_i is input indices[i], x_j is the
          # // j dimension of input x.
          if is_numpy:
              multiplier = input_shape[adv_indx_indices[adv_indx_count - 1]]
              cum_adv_index = tensor_indices[adv_indx_count - 1]
              for i in range(adv_indx_count - 2, -1, -1):
                  adv_index = multiplier * tensor_indices[i]
                  cum_adv_index = cum_adv_index + adv_index
                  multiplier = multiplier * input_shape[adv_indx_indices[i]]
              cum_adv_index = get_trt_tensor(
                  ctx, cum_adv_index, name + "_index_sum_intermediate"
              )
          else:
              multiplier = dim_tensor_list[adv_indx_indices[adv_indx_count - 1]]
              cum_adv_index = tensor_indices[adv_indx_count - 1]
              for i in range(adv_indx_count - 2, -1, -1):
                  adv_index = convert_binary_elementwise(
                      ctx,
                      target,
                      source_ir,
                      name + f"_index_intermediate_{i}",
                      trt.ElementWiseOperation.PROD,
                      multiplier,
                      tensor_indices[i],
                  )
                  cum_adv_index = convert_binary_elementwise(
                      ctx,
                      target,
                      source_ir,
                      name + f"_index_sum_intermediate_{i}",
                      trt.ElementWiseOperation.SUM,
                      cum_adv_index,
                      adv_index,
                  )
                  multiplier = convert_binary_elementwise(
                      ctx,
                      target,
                      source_ir,
                      name + f"_index_intermediate_xj_{i}",
                      trt.ElementWiseOperation.PROD,
                      multiplier,
                      dim_tensor_list[adv_indx_indices[i]],
                  )
  
          gather_layer_element = ctx.net.add_gather(flatten_tensor, cum_adv_index, 0)
          set_layer_name(
              gather_layer_element, target, name + "_index_gather_element", source_ir
          )
          gather_out = gather_layer_element.get_output(0)
          _LOGGER.debug(f"The shape after cumultative gather is {gather_out.shape}")
          _LOGGER.debug(f"The shape for cumulative adv index is {cum_adv_index}")
  
          cum_adv_index_shape_layer = ctx.net.add_shape(cum_adv_index)
          set_layer_name(
              cum_adv_index_shape_layer, target, name + "_cum_adv_index_shape", source_ir
          )
          cum_adv_index_shape_tensor = cum_adv_index_shape_layer.get_output(0)
          cum_adv_index_shape_tensor = cast_trt_tensor(
              ctx,
              cum_adv_index_shape_tensor,
              trt.int32,
              name + "_cum_adv_index_shape_casted",
          )
          cum_adv_index_shape = cum_adv_index.shape
          _LOGGER.debug(f"The shape for cumulative adv index is {cum_adv_index_shape}")
          # check if all advanced indices are consecutive
          concat_tensor_reshape = []
          if (
              adv_indx_count
              == adv_indx_indices[adv_indx_count - 1] - adv_indx_indices[0] + 1
          ):
              _LOGGER.debug("The indices are continuous in this case")
              concat_tensor_reshape.append(
                  get_trt_tensor(ctx, -1, name + "_dynamic_concat")
              )
              for i in range(0, rank):
                  if i not in adv_indx_indices:
                      curr_dim = dim_tensor_list[i]
                      concat_tensor_reshape.append(curr_dim)
  
              concat_tensor_layer = ctx.net.add_concatenation(concat_tensor_reshape)
              set_layer_name(
                  concat_tensor_layer, target, name + "_index_Concat_reshape", source_ir
              )
              concat_tensor = concat_tensor_layer.get_output(0)
  
              regular_index_shuffle_layer = ctx.net.add_shuffle(gather_out)
              regular_index_shuffle_layer.set_input(1, concat_tensor)
              set_layer_name(
                  regular_index_shuffle_layer,
                  target,
                  name + "_index_regular_index",
                  source_ir,
              )
              unfold_tensor = regular_index_shuffle_layer.get_output(0)
              _LOGGER.debug("The tensor is unfolded now")
              _LOGGER.debug(f"The unfolded tensor shape is {unfold_tensor.shape}")
  
              # Transpose folded advanced indexed axis to its original location.
              transpose_advanced_shuffle_layer = ctx.net.add_shuffle(unfold_tensor)
              new_order = []
              for i in range(1, adv_indx_indices[0] + 1):
                  new_order.append(i)
              new_order.append(0)
              for i in range(adv_indx_indices[0] + 1, rank - adv_indx_count + 1):
                  new_order.append(i)
              _LOGGER.debug(f"Transposing the indices to correct position {new_order}")
  
              transpose_advanced_shuffle_layer.second_transpose = tuple(new_order)
              set_layer_name(
                  transpose_advanced_shuffle_layer,
                  target,
                  name + "_index_advanced_shuffle_transpose",
                  source_ir,
              )
              transpose_tensor = transpose_advanced_shuffle_layer.get_output(0)
  
              # unfold advanced layer
              concat_final_tensor = []
              for i in range(0, adv_indx_indices[0]):
                  current_dim = dim_tensor_list[i]
                  concat_final_tensor.append(current_dim)
  
              concat_final_tensor.append(cum_adv_index_shape_tensor)
              for i in range(adv_indx_indices[0], rank):
                  if i not in (adv_indx_indices):
                      current_dim = dim_tensor_list[i]
                      concat_final_tensor.append(current_dim)
  
              concat_final_shape_layer = ctx.net.add_concatenation(concat_final_tensor)
              set_layer_name(
                  concat_final_shape_layer,
                  target,
                  name + "_index_continuous_concat_final_shape_layer",
                  source_ir,
              )
              concat_final_tensor = concat_final_shape_layer.get_output(0)
  
              unfold_advanced_shuffle_layer = ctx.net.add_shuffle(transpose_tensor)
              # check this
              unfold_advanced_shuffle_layer.set_input(1, concat_final_tensor)
              set_layer_name(
                  unfold_advanced_shuffle_layer,
                  target,
                  name + "_unfold_advanced_index",
                  source_ir,
              )
              reshape_output = unfold_advanced_shuffle_layer.get_output(0)
  
          else:
              _LOGGER.debug("The indices are not continuous in this case")
              concat_final_tensor = []
              concat_final_tensor.append(cum_adv_index_shape_tensor)
              for i in range(0, rank):
                  if i not in adv_indx_indices:
                      curr_dim = dim_tensor_list[i]
                      concat_final_tensor.append(curr_dim)
  
              concat_final_shape_layer = ctx.net.add_concatenation(concat_final_tensor)
              set_layer_name(
                  concat_final_shape_layer,
                  target,
                  name + "_index_non_continuous_concat_final_shape_layer",
                  source_ir,
              )
              concat_final_tensor = concat_final_shape_layer.get_output(0)
  
              reshape_layer = ctx.net.add_shuffle(gather_out)
              reshape_layer.set_input(1, concat_final_tensor)
              set_layer_name(
                  reshape_layer,
                  target,
                  name + "_index_non_continuous_shuffle_final_shape_layer",
                  source_ir,
              )
              reshape_output = reshape_layer.get_output(0)
  
          return reshape_output
  
  
  def index_select(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: int,
      index: TRTTensor,
  ) -> TRTTensor:
      # The axis parameter specifies the dimension along which to index.
      dim = get_positive_dim(dim, len(input.shape))
      gather_layer = ctx.net.add_gather(input, index, axis=dim)
  
      set_layer_name(gather_layer, target, f"{name}_gather_layer_default", source_ir)
  
      return gather_layer.get_output(0)
  
  
  def scatter(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: int,
      index: Union[TRTTensor, np.ndarray, torch.Tensor],
      src: Union[TRTTensor, int, float],
  ) -> TRTTensor:
      input_shape = input.shape
      index_shape = index.shape
      index_shape_list = list(index_shape)
      if index.dtype == trt.int64:
          index = cast_trt_tensor(ctx, index, trt.int32, name + "_cast_index_tensor")
      dim = get_positive_dim(dim, len(input_shape))
      src_tensor = src
      # scatter.value
      if isinstance(src, (int, float)):
          src_tensor = get_trt_tensor(
              ctx, src * np.ones(index_shape_list), name + "_value_tensor"
          )
          src_tensor = cast_trt_tensor(
              ctx, src_tensor, input.dtype, name + "_cast_value_tensor"
          )
      # scatter.src
      elif not (isinstance(src, TRTTensor)):
          src_tensor = get_trt_tensor(ctx, src, name + "_src_tensor")
  
      scatter_layer = ctx.net.add_scatter(
          input, index, src_tensor, trt.ScatterMode.ELEMENT
      )
      scatter_layer.axis = dim
      set_layer_name(scatter_layer, target, name + "_scatter_layer", source_ir)
      out = scatter_layer.get_output(0)
      return out
  
  
  def gather(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: int,
      index: Union[TRTTensor, np.ndarray, torch.Tensor],
  ) -> TRTTensor:
      input_shape = input.shape
      dim = get_positive_dim(dim, len(input_shape))
      index = cast_trt_tensor(ctx, index, trt.int32, name + "_cast_index_tensor")
      gather_layer = ctx.net.add_gather(input, index, axis=dim)
      gather_layer.mode = trt.GatherMode.ELEMENT
      set_layer_name(gather_layer, target, name + "_gather_layer_element", source_ir)
      out = gather_layer.get_output(0)
      return out
  
  
  def index_put_converter(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input_tensor: TRTTensor,
      input_indices: Sequence[Union[TRTTensor, torch.Tensor, np.ndarray, int, None]],
      values: TRTTensor,
      accumulate: bool = False,
  ) -> TRTTensor:
      # Convert 'input_indices' to TRT tensors (or keep None as is)
      input_indices = expand_boolean_indices(
          ctx, target, source_ir, name, input_tensor, input_indices
      )
      indices: List[Optional[Union[TRTTensor, None]]] = []
      for i, idx in enumerate(input_indices):
          if idx is None:
              indices.append(None)
          else:
              if not isinstance(idx, TRTTensor):
                  idx = get_trt_tensor(ctx, idx, f"{name}_index_{i}", min_rank=1)
              if len(idx.shape) == 0 or not idx.shape:  # Reshape a scalar to (1,)
                  idx = impl.shuffle.reshape(
                      ctx, target, source_ir, f"{name}_reshape_idx_{i}", idx, (1,)
                  )
              indices.append(idx)
  
      rank = len(input_tensor.shape)
      # Pad the 'indices' list with None for remaining dimensions
      indices = list(indices) + [None] * (rank - len(indices))
  
      # Separate 'F' (Free) dimensions where None is used, and 'I' (Indexed) dimensions
      F = [i for i in range(rank) if indices[i] is None]  # Free dimensions
      I = [i for i in range(rank) if indices[i] is not None]  # Indexed dimensions
      K = len(I)
      # Determine the maximum size 'N' among the index tensors
      if K > 0:
          index_shapes = (
              []
          )  # [tensor.shape[0] for tensor in indices if tensor is not None]
          for idx_tensor in indices:
              if idx_tensor is not None:
                  if idx_tensor.shape[0] != DYNAMIC_DIM:
                      index_shapes.append(idx_tensor.shape[0])
                  else:
                      index_shapes.append(
                          get_shape(
                              ctx,
                              target,
                              source_ir,
                              name + "idx_shape_dim_0",
                              idx_tensor,
                              0,
                          )
                      )
          N = max(index_shapes) if index_shapes else 1
      else:
          N = 1
  
      # Compute shapes and volume for the free dimensions
      F_shapes = [input_tensor.shape[i] for i in F]
      assert -1 not in F_shapes, "Dynamic shape in free dimensions is not supported"
      F_volume = trt.volume(F_shapes) if F_shapes else 1
  
      # Process indexed dimensions (I)
      I_tensors = []
      for i in I:
          idx = indices[i]
          assert idx is not None
          idx_reshaped = impl.unsqueeze.unsqueeze(
              ctx, target, source_ir, f"{name}_unsqueeze_idx_I_{i}", idx, 1
          )
          expanded_idx = impl.slice.expand(
              ctx,
              target,
              source_ir,
              f"{name}_expand_idx_I_{i}",
              idx_reshaped,
              (N, F_volume),
          )
          I_tensors.append(expanded_idx)
  
      # Create a meshgrid for free dimensions (F)
      if len(F) > 0:
          arange_tensors = []
          for dim in F:
              dim_size = input_tensor.shape[dim]
              arange_tensor = impl.arange.arange(
                  ctx, target, source_ir, f"{name}_arange_{dim}", 0, dim_size, 1
              )
              arange_tensors.append(arange_tensor)
  
          if len(arange_tensors) == 1:
              # No need to stack
              meshgrid_stacked = arange_tensors[0]
          else:
              meshgrid_tensors = []
              for i, arange in enumerate(arange_tensors):
                  reshape_shape = [1] * len(F)
                  reshape_shape[i] = F_shapes[i]
                  arange_reshaped = impl.shuffle.reshape(
                      ctx,
                      target,
                      source_ir,
                      f"{name}_reshape_arange_F_{F[i]}",
                      arange,
                      tuple(reshape_shape),
                  )
                  expanded_arange = impl.slice.expand(
                      ctx,
                      target,
                      source_ir,
                      f"{name}_expand_arange_F_{F[i]}",
                      arange_reshaped,
                      tuple(F_shapes),
                  )
                  meshgrid_tensors.append(expanded_arange)
  
              meshgrid_stacked = impl.cat.cat(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_stack_meshgrid",
                  [
                      impl.shuffle.reshape(
                          ctx,
                          target,
                          source_ir,
                          f"{name}_reshape_mesh_{i}",
                          t,
                          (*F_shapes, 1),
                      )
                      for i, t in enumerate(meshgrid_tensors)
                  ],
                  dim=-1,
              )
          meshgrid_reshaped = impl.shuffle.reshape(
              ctx,
              target,
              source_ir,
              f"{name}_reshape_meshgrid",
              meshgrid_stacked,
              (F_volume, len(F)),
          )
          if K > 0:
              meshgrid_expanded = impl.slice.expand(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_expand_meshgrid",
                  meshgrid_reshaped,
                  (N, F_volume, len(F)),
              )
          else:
              meshgrid_expanded = meshgrid_reshaped
      else:
          meshgrid_expanded = None
  
      # Combine all indexed dimensions (I)
      if K > 0:
  
          I_combined = [
              impl.shuffle.reshape(
                  ctx, target, source_ir, f"{name}_reshape_I_{i}", t, (N, F_volume, 1)
              )
              for i, t in enumerate(I_tensors)
          ]
      else:
          I_combined = []
  
      # Build the final index list (ii_list) by slicing either I_combined or meshgrid_expanded
      ii_list = []
      i_idx = 0
      f_idx = 0
      for dim in range(rank):
          unique_suffix = f"{dim}_{i_idx if dim in I else f_idx}"
          if dim in I:
              idx_tensor = I_combined[i_idx]
              ii_list.append(idx_tensor)
              i_idx += 1
          else:
              start = [0, 0, f_idx]
              shape = [-1, F_volume, 1] if isinstance(N, TRTTensor) else [N, F_volume, 1]
              stride = [1, 1, 1]
              mesh_tensor = impl.slice.slice(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_slice_F_dim_{unique_suffix}",
                  meshgrid_expanded,
                  start,
                  shape,
                  stride,
              )
              ii_list.append(mesh_tensor)
              f_idx += 1
  
      # Concatenate the final indices and reshape to (N * F_volume, rank)
      indices_cat = impl.cat.cat(
          ctx, target, source_ir, f"{name}_cat_indices", ii_list, dim=2
      )
  
      # Flatten the indices_cat to (N * F_volume, rank)
      indices_cat = impl.shuffle.reshape(
          ctx,
          target,
          source_ir,
          f"{name}_reshape_indices_cat",
          indices_cat,
          (-1, rank),
      )
  
      if not isinstance(values, TRTTensor):
          values = get_trt_tensor(ctx, values, f"{name}_values", min_rank=0)
  
      # Define the expected shape based on (N,) + F_shapes
      expected_shape = (
          (-1,) + tuple(F_shapes) if isinstance(N, TRTTensor) else (N,) + tuple(F_shapes)
      )
  
      # Broadcast 'values' to match the expected shape
      if len(values.shape) == 0 or values.shape == (1,):  # Scalar case
          values_reshaped = impl.shuffle.reshape(
              ctx, target, source_ir, f"{name}_reshape_scalar", values, (1,)
          )
          values_expanded = impl.slice.expand(
              ctx,
              target,
              source_ir,
              f"{name}_expand_values_scalar",
              values_reshaped,
              expected_shape,
          )
      else:  # Non-scalar case
          values_shape = list(values.shape)
          if (
              K > 0
              and N in values_shape
              and (len(F) > 1 and max(F) - min(F) + 1 == len(F))
          ):
              # Continuous case
              n_idx = values_shape.index(N)
              permute_order = [n_idx] + [
                  i for i in range(len(values_shape)) if i != n_idx
              ]
              values_permuted = impl.permutation.permute(
                  ctx, target, source_ir, f"{name}_permute_values", values, permute_order
              )
              remaining_shape = [
                  values_shape[i] for i in range(len(values_shape)) if i != n_idx
              ]
              target_f_dims = len(F)
              current_f_dims = len(remaining_shape)
              if current_f_dims < target_f_dims:
                  values_expanded_shape = (
                      [N] + [1] * (target_f_dims - current_f_dims) + remaining_shape
                  )
              else:
                  values_expanded_shape = [N] + remaining_shape[:target_f_dims]
              values_expanded = impl.shuffle.reshape(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_unsqueeze_values",
                  values_permuted,
                  tuple(values_expanded_shape),
              )
              broadcast_shape = []
              for exp_dim, val_dim in zip(expected_shape, values_expanded_shape):
                  if val_dim == 1:
                      broadcast_shape.append(exp_dim)
                  elif val_dim == exp_dim:
                      broadcast_shape.append(val_dim)
                  else:
                      raise ValueError(
                          f"Cannot broadcast {values_expanded_shape} to {expected_shape}"
                      )
              values_expanded = impl.slice.expand(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_expand_values",
                  values_expanded,
                  tuple(broadcast_shape),
              )
          else:
              # Discontinuous case
              values_shape_padded = [1] * (
                  len(expected_shape) - len(values.shape)
              ) + list(values.shape)
              broadcast_shape = []
              for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
                  if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM:
                      broadcast_shape.append(-1)
                  elif val_dim == 1 or exp_dim == val_dim:
                      broadcast_shape.append(exp_dim)
                  else:
                      raise ValueError(
                          f"Cannot broadcast {values.shape} to {expected_shape}"
                      )
  
              values_expanded = impl.slice.expand(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_expand_values",
                  values,
                  expected_shape,
              )
  
      # Flatten values to (N * F_volume,)
      flattened_values = impl.shuffle.reshape(
          ctx,
          target,
          source_ir,
          f"{name}_flatten_values",
          values_expanded,
          (-1,),
      )
      indices_cat = cast_trt_tensor(ctx, indices_cat, trt.int32, f"{name}_idx_int32")
      if accumulate:
          zero_tensor = impl.full.full(
              ctx,
              target,
              source_ir,
              f"{name}_zero_tensor",
              [
                  get_shape(
                      ctx,
                      target,
                      source_ir,
                      name + f"input_tensor_shape_dim_{i}",
                      input_tensor,
                      i,
                  )
                  for i in range(len(input_tensor.shape))
              ],
              0.0,
              dtype=input_tensor.dtype,
          )
          # Perform Scatter ND operation
          scatter_layer = ctx.net.add_scatter(
              zero_tensor,
              indices_cat,
              flattened_values,
              trt.ScatterMode.ND,
          )
          set_layer_name(scatter_layer, target, f"{name}_scatter", source_ir)
  
          scatter_out = scatter_layer.get_output(0)
          result = impl.elementwise.add(
              ctx, target, source_ir, f"{name}_add", scatter_out, input_tensor
          )
          return result
  
      else:
          scatter_layer = ctx.net.add_scatter(
              input_tensor,
  