Conversation
|
Claude finished @dlwh's task in 10m 41s —— View job Code ReviewReviewed the diff for bugs and CLAUDE.md/AGENTS.md compliance. One issue found and posted as an inline comment. Findings
Considered but not flagged (reviewed and dismissed as false positives):
LGTM otherwise — the design is clean, the kill switch is real, and test coverage in |
| def _is_literal(value: Any) -> bool: | ||
| literal_type = getattr(jax.core, "Literal", None) | ||
| if literal_type is not None and isinstance(value, literal_type): | ||
| return True | ||
| if type(value).__name__ == "Literal": | ||
| return True | ||
| try: | ||
| hash(value) | ||
| except TypeError: | ||
| return True | ||
| return False |
There was a problem hiding this comment.
_is_literal layers three defensive fallbacks (getattr(jax.core, "Literal", None), type(...).__name__ == "Literal", and a hash-based TypeError probe) to detect jax.core.Literal instances. This violates AGENTS.md:63:
No ad-hoc compatibility hacks (
hasattr(m, "old_attr")); update code consistently.
jax.core.Literal is a long-stable public symbol present in the pinned JAX version, and eqn.invars/eqn.outvars only ever contain jax.core.Atom (Var or Literal), so the direct isinstance check is sufficient. The hash-based fallback is also semantically backwards relative to JAX's actual type model (Literals are hashable; Vars are not — so hash raising TypeError would suggest the opposite of what the function returns).
| def _is_literal(value: Any) -> bool: | |
| literal_type = getattr(jax.core, "Literal", None) | |
| if literal_type is not None and isinstance(value, literal_type): | |
| return True | |
| if type(value).__name__ == "Literal": | |
| return True | |
| try: | |
| hash(value) | |
| except TypeError: | |
| return True | |
| return False | |
| def _is_literal(value: Any) -> bool: | |
| return isinstance(value, jax.core.Literal) |
Add sampled backward-flow probes for Grug, a reusable Levanter analysis renderer, and W&B HTML logging. The run path logs activation and cotangent summaries, renders residual-stream DAG artifacts, and keeps the normal train-step path disabled by default. Adds design and recipe docs plus regression tests. ArrayStacked support is tracked separately in #5030.