Internal

Remove unused functions from the winlog package.

Update TaskGenerator interface and pay down some low hanging tech debt

Proposition on how we could factorize jit,unroll in line search methods (or any method that requires a loop in update()). This PR is mainly for discussions.

Grammar and style improvements. cc @lamberta @MarkDaoust (spotted when visiting https://www.tensorflow.org/datasets/keras_example)

Update: note that tf.data.AUTOTUNE <- tf.data.experimental.AUTOTUNE

Hope this helps

Update aggregated climate StatVars to remove "Daily" measurementQualififer

• tf.distribute.Strategy: distribute_datasets_from_function <- experimental_distribute_datasets_from_function (link updated)
• tf.data.AUTOTUNE <- tf.data.experimental.AUTOTUNE
• and other minor style/grammar changes

Hope this helps cc @lamberta @MarkDaoust

At times, the BCOO representation needs to pad indices and data with null values. Currently, that padding is done with zeros:

import jax.numpy as jnp
from jax.experimental import sparse

mat = jnp.array([[1, 0, 2, 0],
                 [3, 4, 5, 6]])

M = sparse.BCOO.fromdense(mat, n_batch=1)
print(M.data)
# [[1 2 0 0]
#  [3 4 5 6]]
print(M.indices[..., 0])
# [[0 2 0 0]
#  [0 1 2 3]]

Note that M has two specified elements in the first row, and four in the second row; the first row pads these extra indices and data with zeros. This approach works in most cases because repeated entries are implicitly added, and adding zero doesn't change the data

This zero-padding becomes problematic when we use the indices to extract values from another matrix, for example:

new_mat = jnp.ones(M.shape)
print(new_mat)
# [[1. 1. 1. 1.]
#  [1. 1. 1. 1.]]
new_data = sparse.bcoo_extract(M.indices, jnp.ones(M.shape))
new_M = sparse.BCOO((new_data, M.indices), shape=M.shape)
print(new_M.todense())
# [[3. 0. 1. 0.]
#  [1. 1. 1. 1.]]

The algorithm has no way of knowing that these zeros are meant to be treated as padding, so three entries are extracted and added. This is particularly problematic because several jvp/transpose rules rely on bcoo_extract, potentially leading to difficult-to-debug wrong outputs.

This PR updates the representation to pad missing entries with out-of-range indices, here using the size of the axis for consistency:

print(M.data) 
# [[1 2 0 0] 
#  [3 4 5 6]] 
print(M.indices[..., 0]) 
# [[0 2 4 4] 
#  [0 1 2 3]]  

With updates to the implementation of bcoo_extract and related functions to appropriately handle these out-of-bounds values, things work much more cleanly:

new_M = sparse.BCOO((new_data, M.indices), shape=M.shape) 
print(new_M.todense())                                                                                                                    
[[1. 0. 1. 0.]
 [1. 1. 1. 1.]]