<commit_before># `jax.debug.print` and `jax.debug.breakpoint`

The {mod}`jax.debug` package offers some useful tools for inspecting values
inside of JIT-ted functions.

## Debugging with `jax.debug.print` and other debugging callbacks

**TL;DR** Use {func}`jax.debug.print` to print values to stdout in `jit`-,`pmap`-, and `pjit`-decorated functions:

```python
import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  jax.debug.print("🤯 {x} 🤯", x=x)
  y = jnp.sin(x)
  jax.debug.print("🤯 {y} 🤯", y=y)
  return y
  
f(2.)
# Prints:
# 🤯 2.0 🤯
# 🤯 0.9092974662780762 🤯
```

<!-- mattjj added this line -->
With some transformations, like `jax.grad` and `jax.vmap`, you can use Python's builtin `print` function to print out numerical values. But `print` won't work with `jax.jit` or `jax.pmap` because those transformations delay numerical evaluation. So use `jax.debug.print` instead!

Semantically, `jax.debug.print` is roughly equivalent to the following Python function

```python
def debug.print(fmt: str, *args: PyTree[Array], **kwargs: PyTree[Array]) -> None:
  print(fmt.format(*args, **kwargs))
```
except that it can be staged out and transformed by JAX. See the {func}`API reference <jax.debug.print>` for more details.

### Why "_debug_" print?
In the name of debugging, `jax.debug.print` can reveal information about _how_ computations are evaluated:

```python
xs = jnp.arange(3.)

def f(x):
  jax.debug.print("x: {}", x)
  y = np.sin(x)
  jax.debug.print("y: {}", y)
  return y
jax.vmap(f)(xs)
# Prints: x: 0.0
#         x: 1.0
#         x: 2.0
#         y: 0.0
#         y: 0.841471
#         y: 0.9092974
jax.lax.map(f, xs)
# Prints: x: 0.0
#         y: 0.0
#         x: 1.0
#         y: 0.841471
#         x: 2.0
#         y: 0.9092974
```
Notice that the printed results are in different orders!

By revealing these inner-workings, the output of `jax.debug.print` doesn't respect JAX's usual semantics guarantees, like that `jax.vmap(f)(xs)` and `jax.lax.map(f, xs)` compute the same thing (in different ways). Yet these evaluation order details are exactly what we might want to see when debugging!

So use `jax.debug.print` for debugging, and not when semantics guarantees are important.
<commit_msg>Fix typo in `jax.debug.print()` documentation.<commit_after># `jax.debug.print` and `jax.debug.breakpoint`

The {mod}`jax.debug` package offers some useful tools for inspecting values
inside of JIT-ted functions.

## Debugging with `jax.debug.print` and other debugging callbacks

**TL;DR** Use {func}`jax.debug.print` to print values to stdout in `jit`-,`pmap`-, and `pjit`-decorated functions:

```python
import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  jax.debug.print("🤯 {x} 🤯", x=x)
  y = jnp.sin(x)
  jax.debug.print("🤯 {y} 🤯", y=y)
  return y
  
f(2.)
# Prints:
# 🤯 2.0 🤯
# 🤯 0.9092974662780762 🤯
```

<!-- mattjj added this line -->
With some transformations, like `jax.grad` and `jax.vmap`, you can use Python's builtin `print` function to print out numerical values. But `print` won't work with `jax.jit` or `jax.pmap` because those transformations delay numerical evaluation. So use `jax.debug.print` instead!

Semantically, `jax.debug.print` is roughly equivalent to the following Python function

```python
def debug.print(fmt: str, *args: PyTree[Array], **kwargs: PyTree[Array]) -> None:
  print(fmt.format(*args, **kwargs))
```
except that it can be staged out and transformed by JAX. See the {func}`API reference <jax.debug.print>` for more details.

### Why "_debug_" print?
In the name of debugging, `jax.debug.print` can reveal information about _how_ computations are evaluated:

```python
xs = jnp.arange(3.)

def f(x):
  jax.debug.print("x: {}", x)
  y = jnp.sin(x)
  jax.debug.print("y: {}", y)
  return y
jax.vmap(f)(xs)
# Prints: x: 0.0
#         x: 1.0
#         x: 2.0
#         y: 0.0
#         y: 0.841471
#         y: 0.9092974
jax.lax.map(f, xs)
# Prints: x: 0.0
#         y: 0.0
#         x: 1.0
#         y: 0.841471
#         x: 2.0
#         y: 0.9092974
```
Notice that the printed results are in different orders!

By revealing these inner-workings, the output of `jax.debug.print` doesn't respect JAX's usual semantics guarantees, like that `jax.vmap(f)(xs)` and `jax.lax.map(f, xs)` compute the same thing (in different ways). Yet these evaluation order details are exactly what we might want to see when debugging!

So use `jax.debug.print` for debugging, and not when semantics guarantees are important.

{'fraction_non_alphanumeric': 0.14244186046511628, 'fraction_numerical': 0.02996422182468694, 'mean_word_length': 3.5090725806451615, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2117558', 'n_tokens_mistral': 1736, 'n_tokens_neox': 1568, 'n_words': 546}