Functionality Roadmap

[seanmor5]

An issue to track some baseline functionality:

Activations

• celu
• elu
• exp
• gelu
• hard_tanh
• hard_sigmoid
• hard_silu/hard_swish
• leaky_relu
• log_sigmoid
• relu
• relu6
• selu
• sigmoid
• silu
• softmax
• softplus
• softsign
• tanh

Initializers

• glorot_uniform
• glorot_normal
• he_normal
• he_uniform
• lecun_uniform
• lecun_normal
• normal
• ones
• orthogonal - requires Support vmap nx#174
• uniform
• zeros

Loss Functions

• binary_crossentropy
• categorical_crossentropy
• categorical_hinge
• cosine_similarity - requires Support vmap nx#174
• ctc
• hinge
• kl_divergence
• log_cosh
• margin_ranking
• mean_absolute_error
• mean_squared_error
• poisson
• soft_margin

Metrics

• accuracy
• mean_squared_error - requires defndelegate
• mean_absolute_error - requires defndelegate
• precision
• recall
• sensitivity
• specificty

Optimizers

Optax style transformations:

• scale
• scale_by_adam
• scale_by_rss
• scale_by_belief
• scale_by_rms
• trace
• clip
• clip_by_global_norm
• centralize
• scale_by_trust_ratio
• scale_by_schedule
• scale_by_radam
• scale_by_stddev

Schedules

• polynomial_schedule
• exponential_decay_schedule
• cosine_decay_schedule
• constant_schedule

Layers

For now, just functional implementations resembling torch.nn.functional or tf.nn:

Linear Layers

• dense
• bilinear - requires Make Nx.dot/4 aware of batch dimensions nx#182

Convolutional Layers

• conv
• conv_transpose
• depthwise_conv
• separable_conv2d
• separable_conv3d

Pooling Layers

• avg_pool
• max_pool
• lp_pool
• adaptive_avg_pool
• adaptive_max_pool
• adaptive_lp_pool
• global_avg_pool
• global_max_pool
• global_lp_pool

Normalization Layers

• batch_norm
• group_norm
• instance_norm
• layer_norm

Dropout Layers

• dropout
• alpha_dropout
• feature_alpha_dropout
• spatial_dropout

Attention Layers

• dot_product_attention - requires Make Nx.dot/4 aware of batch dimensions nx#182
• additive_attention - requires repeat/gather on Nx

Visual Layers

• resize

We can drop off dimensional suffixes in favor of generic implementations too.

[seanmor5]

Closing in favor of smaller issues