Computes a mask over sequence positions for each given length. Args: lengths: [...]. int32 max_len: T, int dtype: outputs dtype. Returns: Tensor [..., T]. 1 is valid and 0 is padding.
(*, lengths: Tensor, max_len: int, dtype: jnp.dtype = jnp.bool)
| 2021 | f"Unable to infer -1 in mesh shape {mesh_shape} as num_devices {num_devices} " |
| 2022 | f"is not a multiple of the product {prod} of mesh axes." |
| 2023 | ) |
| 2024 | |
| 2025 | return tuple(x if x != -1 else num_devices // prod for x in mesh_shape) |
| 2026 | |
| 2027 | |
| 2028 | def thread_stack_traces() -> Sequence[Sequence[str]]: |
| 2029 | """Retrieves the current python stack traces.""" |
| 2030 | grouped_lines = [] |
| 2031 | for thread in threading.enumerate(): |
| 2032 | lines = [] |
| 2033 | thread_id = thread.ident |
| 2034 | lines.append(f"Thread: {thread.name}({thread_id})") |
| 2035 | # pylint: disable-next=protected-access |
| 2036 | for line in traceback.format_stack(sys._current_frames()[thread_id]): |
| 2037 | lines.append(f">>> {line.rstrip()}") |
| 2038 | grouped_lines.append(lines) |
| 2039 | return grouped_lines |
| 2040 | |
| 2041 | |
| 2042 | def pytree_children(node: Any) -> Sequence[tuple[KeyEntry, Any]]: |