code: def __init__(
      self,
      index_map: Callable[..., Any] | None = None,
      block_shape: tuple[int | None, ...] | None = None,
      memory_space: Any = None,
  ):
      self.index_map = index_map
      if block_shape is not None and not isinstance(block_shape, tuple):
          block_shape = tuple(block_shape)
      self.block_shape = block_shape
      self.memory_space = memory_space
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import functools
  from typing import Any, Callable
  from collections.abc import Iterator
  from jax._src import api_util
  from jax._src import core
  from jax._src import linear_util
  from jax._src import state
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import partial_eval
  from jax._src.state import discharge
  import jax.numpy
label_desc: Use image node_id 1 to create a new BlockSpec object  with example: obj = BlockSpec(index_map, block_shape, memory_space)