How deeply does Deep Learning with JAX cover naive approach?

Explains direct application, manual looping, silent shape errors, and performance tradeoffs, but defers efficient vectorization internals.

Naive approach

As explained in Deep Learning with JAX

A simple but potentially inefficient way to apply a scalar function to arrays, either by relying on accidental broadcasting or by looping over elements manually.

A naive approach applies a scalar function to arrays without deliberately arranging vectorization. It may work through broadcasting and universal functions, fail clearly, or silently produce an incorrect result. section 6.1.1

Two straightforward forms illustrate the tradeoff:

Pass arrays directly to the original function. For ten pairs of vectors, dot returns one scalar instead of ten pairwise dot products; broadcasting can let the mistake spread unnoticed. section 6.1.1
Loop over elements manually with a list comprehension. This produces the expected individual results, but it is straightforward rather than efficient. section 6.1.1

Listing 6.7 Comparing the speed of different approaches

%timeit [dot(v1s[i],v2s[i]).block_until_ready() for i in range(v1s.shape[0])]                               #1

>>> 4.58 ms ± 152 µs per loop                               #2
(mean ± std. dev. of 7 runs, 100 loops each)              #2
 
%timeit dot_vectorized(v1s, v2s).block_until_ready()
 
>>> 91.9 µs ± 20 µs per loop                                #3
(mean ± std. dev. of 7 runs, 10000 loops each)            #3
 
%timeit dot_vmapped(v1s, v2s).block_until_ready()
 
>>> 831 µs ± 7.7 µs per loop                                #4
(mean ± std. dev. of 7 runs, 1000 loops each)             #4
 
dot_vectorized_jitted = jax.jit(dot_vectorized)             #5
dot_vmapped_jitted = jax.jit(dot_vmapped)                   #5
 
# warm-up
dot_vectorized_jitted(v1s, v2s);
dot_vmapped_jitted(v1s, v2s);
 
%timeit dot_vectorized_jitted(v1s, v2s).block_until_ready()
 
>>> 6.93 µs ± 194 ns per loop                                #6
(mean ± std. dev. of 7 runs, 100000 loops each)
 
%timeit dot_vmapped_jitted(v1s, v2s).block_until_ready()
 
>>> 7.6 µs ± 1.42 µs per loop    
(mean ± std. dev. of 7 runs, 100000 loops each)            #6
#1 Don’t forget about asynchronous dispatch!
#2 The naive approach is the slowest one.
#3 The manually vectorized function has great performance.
#4 The automatically vectorized function is slower but still very fast.
#5 JIT-ting both vectorized functions
#6 After JIT-ting, the speeds of manually and automatically vectorized solutions are much closer.

Speed measurements show the naive loop is much slower than vectorized alternatives; JIT compilation makes the approaches more similar. section 6.1.4

How Deep Learning with JAX covers it

liveBook manning.com

Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

algorithmcontrasts with · An algorithm is a finite, precisely specified procedure for performing a computation, transforming inputs into outputs, or solving a problem.
memoizationalternative to · Memoization is an optimization technique that stores computation results with their inputs so equivalent later evaluations can reuse them without recomputation.
tail recursionalternative to · Tail recursion is recursion in which the recursive call is the final operation of the current function invocation.
vectorizationcontrasts with · Vectorization is the representation of data or computations as vectors or arrays so they can be processed numerically or collectively.
inner productuses · An inner product is a scalar-valued operation on two vectors that combines their corresponding components, summing their products and, for complex vectors, using conjugation.
brute-force searchincludes · Brute-force search systematically tests many or all candidates in a defined search space, without relying on selective shortcuts to narrow it.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.