code: def _next_power_of_2_or_max(n: int, max_n: int) -> int:
      """Return the smallest power of 2 greater than or equal to n, with a limit.
  
      Useful when used in splitting a tensor into chunks with power-of-2 sizes.
      """
      # special case, just split to 1 element chunks.
      if n == 0:
          return 1
      orig_n = n
      n -= 1
      n |= n >> 1
      n |= n >> 2
      n |= n >> 4
      n |= n >> 8
      n |= n >> 16
      n += 1
      assert n >= orig_n, f"{n} vs. {orig_n}"
      assert bin(n).count("1") == 1, bin(
          n
      )  # Catch the case n is too large for this function.
      if n > max_n:
          return max_n
      return n
  
imports: from typing import Any, Optional, Tuple
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: _next_power_of_2_or_max(n, max_n) and returns: n, int, max_n