jax.extend.pallas.GridMapping
=============================

.. currentmodule:: jax.extend.pallas

.. autoclass:: GridMapping

   
   .. automethod:: __init__

   
   .. rubric:: Methods

   .. autosummary::
   
      ~GridMapping.__init__
      ~GridMapping.check_invariants
      ~GridMapping.replace
      ~GridMapping.to_lojax
      ~GridMapping.trace_env
   
   

   
   
   .. rubric:: Attributes

   .. autosummary::
   
      ~GridMapping.block_mappings_output
      ~GridMapping.debug
      ~GridMapping.get_grid_indices
      ~GridMapping.in_shapes
      ~GridMapping.local_grid_env
      ~GridMapping.num_dynamic_grid_bounds
      ~GridMapping.num_scratch_operands
      ~GridMapping.out_shapes
      ~GridMapping.slice_block_ops
      ~GridMapping.slice_index_ops
      ~GridMapping.slice_scratch_ops
      ~GridMapping.static_grid
      ~GridMapping.grid
      ~GridMapping.grid_names
      ~GridMapping.block_mappings
      ~GridMapping.index_map_tree
      ~GridMapping.index_map_avals
      ~GridMapping.vmapped_dims
      ~GridMapping.scratch_avals
      ~GridMapping.num_index_operands
      ~GridMapping.num_inputs
      ~GridMapping.num_outputs