More fixes for arrays with large sizes (#1405)

* compile works for big arrays when contiguous

* style

* nits in docs

* a bunch more stuff

* update jit

* update jit

* use constant for shapes and strides and remove elem_to_loc overload

* use kernel instantiation

* docs nits

* update binary and ternary

* comments

docs/src/python/nn/functions.rst

@@ -13,6 +13,7 @@ simple functions.
    :template: nn-module-template.rst
 
    elu
+   celu
    gelu
    gelu_approx
    gelu_fast_approx

docs/src/python/nn/layers.rst

@@ -13,6 +13,7 @@ Layers
    AvgPool1d
    AvgPool2d
    BatchNorm
+   CELU
    Conv1d
    Conv2d
    Conv3d
@@ -23,6 +24,7 @@ Layers
    Dropout2d
    Dropout3d
    Embedding
+   ELU
    GELU
    GLU
    GroupNorm
@@ -34,6 +36,8 @@ Layers
    LayerNorm
    LeakyReLU
    Linear
+   LogSigmoid
+   LogSoftmax
    LSTM
    MaxPool1d
    MaxPool2d
@@ -49,6 +53,7 @@ Layers
    RoPE
    SELU
    Sequential
+   Sigmoid
    SiLU
    SinusoidalPositionalEncoding
    Softmin