🤔 Config

Defaults:

Type Default Details
precision int 3 Digits after .
threshold_max int 3 .abs() larger than 1e3 -> Sci mode
threshold_min int -4 .abs() smaller that 1e-4 -> Sci mode
sci_mode NoneType None Sci mode (2.3e4). None=auto
show_mem_above int 1024 Show memory footprint above this size
indent int 2 Indent for .deeper()
color bool True ANSI colors in text
deeper_width int 9 For .deeper, width per level
plt_seed int 42 Sampling seed for plot
fig_close bool True Close matplotlib Figure
fig_show bool False Call plt.show() for .plt, .chans and .rgb

source

set_config

 set_config (precision:Union[~Default,int,NoneType]=Ignore,
             threshold_min:Union[~Default,int,NoneType]=Ignore,
             threshold_max:Union[~Default,int,NoneType]=Ignore,
             sci_mode:Union[~Default,bool,NoneType]=Ignore,
             show_mem_above:Union[~Default,int,NoneType]=Ignore,
             indent:Union[~Default,bool,NoneType]=Ignore,
             color:Union[~Default,bool,NoneType]=Ignore,
             deeper_width:Union[~Default,int,NoneType]=Ignore,
             plt_seed:Union[~Default,int,NoneType]=Ignore,
             fig_close:Union[~Default,bool,NoneType]=Ignore,
             fig_show:Union[~Default,bool,NoneType]=Ignore)

Set config variables

source

get_config

 get_config ()

Get a copy of config variables

source

config

 config (precision:Union[~Default,int,NoneType]=Ignore,
         threshold_min:Union[~Default,int,NoneType]=Ignore,
         threshold_max:Union[~Default,int,NoneType]=Ignore,
         sci_mode:Union[~Default,bool,NoneType]=Ignore,
         show_mem_above:Union[~Default,int,NoneType]=Ignore,
         indent:Union[~Default,bool,NoneType]=Ignore,
         color:Union[~Default,bool,NoneType]=Ignore,
         deeper_width:Union[~Default,int,NoneType]=Ignore,
         plt_seed:Union[~Default,int,NoneType]=Ignore,
         fig_close:Union[~Default,bool,NoneType]=Ignore,
         fig_show:Union[~Default,bool,NoneType]=Ignore)

Context manager for temporarily setting printting options.

Examples

import jax.numpy as jnp
from lovely_jax import set_config, config, monkey_patch
monkey_patch()

Precision

set_config(precision=5)
jnp.array([1., 2, jnp.nan])
Array[3] μ=1.50000 σ=0.50000 NaN! cpu:0 [1.00000, 2.00000, nan]
jnp.array([1., 2, jnp.nan])
Array[3] μ=1.50000 σ=0.50000 NaN! cpu:0 [1.00000, 2.00000, nan]

Scientific mode

set_config(sci_mode=True) # Force always on
str(jnp.array([1., 2, jnp.nan]))
'Array[3] μ=1.50000e+00 σ=5.00000e-01 \x1b[31mNaN!\x1b[0m cpu:0 [1.00000e+00, 2.00000e+00, nan]'

Color on/off

set_config(color=False) # Force always off
jnp.array([1., 2, jnp.nan])
Array[3] μ=1.50000e+00 σ=5.00000e-01 NaN! cpu:0 [1.00000e+00, 2.00000e+00, nan]
test_array_repr(str(jnp.array([1., 2, jnp.nan])),
        'Array[3] μ=1.50000e+00 σ=5.00000e-01 NaN! gpu:0 [1.00000e+00, 2.00000e+00, nan]')

Control .deeper

set_config(deeper_width=3)
image = jnp.load("mysteryman.npy")
image = image.at[1,100,100].set(jnp.nan)

image.deeper(2)
Array[3, 196, 196] n=115248 (0.4Mb) x∈[-2.11790e+00, 2.64000e+00] μ=-3.88310e-01 σ=1.07318e+00 NaN! cpu:0
  Array[196, 196] n=38416 x∈[-2.11790e+00, 2.24891e+00] μ=-3.24352e-01 σ=1.03586e+00 cpu:0
    Array[196] x∈[-1.91241e+00, 2.24891e+00] μ=-6.73483e-01 σ=5.20629e-01 cpu:0
    Array[196] x∈[-1.86103e+00, 2.16328e+00] μ=-7.38488e-01 σ=4.17012e-01 cpu:0
    Array[196] x∈[-1.75828e+00, 2.19753e+00] μ=-8.05501e-01 σ=3.95835e-01 cpu:0
    ...
  Array[196, 196] n=38416 x∈[-1.96569e+00, 2.42857e+00] μ=-2.73903e-01 σ=9.72652e-01 NaN! cpu:0
    Array[196] x∈[-1.86064e+00, 2.41106e+00] μ=-5.28772e-01 σ=5.54540e-01 cpu:0
    Array[196] x∈[-1.82563e+00, 2.35854e+00] μ=-5.61732e-01 σ=4.71564e-01 cpu:0
    Array[196] x∈[-1.75560e+00, 2.37605e+00] μ=-6.21756e-01 σ=4.57265e-01 cpu:0
    ...
  Array[196, 196] n=38416 x∈[-1.80444e+00, 2.64000e+00] μ=-5.66674e-01 σ=1.17775e+00 cpu:0
    Array[196] x∈[-1.71730e+00, 2.39599e+00] μ=-9.81537e-01 σ=3.49106e-01 cpu:0
    Array[196] x∈[-1.75216e+00, 2.32627e+00] μ=-1.03418e+00 σ=3.13168e-01 cpu:0
    Array[196] x∈[-1.64758e+00, 2.37856e+00] μ=-1.08647e+00 σ=3.13411e-01 cpu:0
    ...
test_eq(len(str(image.deeper(2))), 1127)

Reser to defaults

set_config(precision=None, sci_mode=None, color=None, deeper_width=None)
str(jnp.array([1., 2, jnp.nan]))
'Array[3] μ=1.500 σ=0.500 \x1b[31mNaN!\x1b[0m cpu:0 [1.000, 2.000, nan]'
test_array_repr(str(jnp.array([1., 2, jnp.nan])),
    'Array[3] μ=1.500 σ=0.500 \x1b[31mNaN!\x1b[0m gpu:0 [1.000, 2.000, nan]')

Context manager

display(jnp.array([1., 2, jnp.nan]))
with config(sci_mode=True, color=False):
    display(jnp.array([1., 2, jnp.nan]))
display(jnp.array([1., 2, jnp.nan]))
Array[3] μ=1.500 σ=0.500 NaN! cpu:0 [1.000, 2.000, nan]
Array[3] μ=1.500e+00 σ=5.000e-01 NaN! cpu:0 [1.000e+00, 2.000e+00, nan]
Array[3] μ=1.500 σ=0.500 NaN! cpu:0 [1.000, 2.000, nan]

Matplotlib and seed

# torch.manual_seed(42)
# a = torch.randn(1000)

key = jax.random.PRNGKey(0)
a = jax.random.normal(key, (1000,))
_ = a.plt() # The figure was closed, nothing is displayed
set_config(fig_close=False)
_ = a.plt() # figure was not closed. All figures that are not closed are displayed after the cell runs.

For performance reasons, .plt will randomly sample up tp max_s elements from the data (10k be default).

You can change the seed used for this sampling (42 by default):

set_config(plt_seed=1)
a.plt(max_s=100)

set_config(plt_seed=2)
a.plt(max_s=100)

More details in matplotlib