docstr: Allow indexing when leading batch / time dimensions are
  present in the parameters
  
  !!! example
      ```python
      ...
      domain = Domain((16, (1.0,))
  
      # 10 fields
      params = random.uniform(key, (10, 16, 1))
      a = OnGrid(params, domain)
  
      # Field at the 5th index
      field = a[5]
      ```
  
  Returns:
    OnGrid: A linear discretization on the grid points of the domain.
  
  Raises:
    IndexError: If the field is not indexable (single field).
docstr_tok: ["Allow","indexing","when","leading","batch","\/","time","dimensions","are","present","in","the","parameters","!","!","!","example","``","`","python","...","domain","=","Domain","(","(","16",",","(","1.0",",",")",")","#","10","fields","params","=","random.uniform","(","key",",","(","10",",","16",",","1",")",")","a","=","OnGrid","(","params",",","domain",")","#","Field","at","the","5th","index","field","=","a","[","5","]","``","`","Returns",":","OnGrid",":","A","linear","discretization","on","the","grid","points","of","the","domain",".","Raises",":","IndexError",":","If","the","field","is","not","indexable","(","single","field",")","."]
code: def __getitem__(self, idx):
      r"""Allow indexing when leading batch / time dimensions are
      present in the parameters
  
      !!! example
          ```python
          ...
          domain = Domain((16, (1.0,))
  
          # 10 fields
          params = random.uniform(key, (10, 16, 1))
          a = OnGrid(params, domain)
  
          # Field at the 5th index
          field = a[5]
          ```
  
      Returns:
        OnGrid: A linear discretization on the grid points of the domain.
  
      Raises:
        IndexError: If the field is not indexable (single field).
      """
      if self.domain.ndim + 1 == len(self.params.shape):
          raise IndexError(
              "Indexing is only supported if there's at least one batch / time dimension"
          )
  
      new_params = self.params[idx]
      return self.__class__(new_params, self.domain)
  
imports: import warnings
  from typing import Callable, TypeVar
  import equinox
  from equinox import tree_equal
  from jax import eval_shape
  from jax import numpy
  from jax import vmap
  from jaxtyping import PyTree
  from jaxdf.core import Field
  from jaxdf.geometry import Domain
label_desc: Use image node_id 4 for calling the OnGrid obj's underlying member method code with example usage: obj.__getitem__(idx) and returns: self