feat: add bf16 function (#1104)

* feat: add bf16 function

* docs: add bf16 to docs

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.3.4"
+version = "1.4.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api/Lux/utilities.md

@@ -89,6 +89,7 @@ simply pass the parameters / states / arrays into one of the following functions
 Lux.f16
 Lux.f32
 Lux.f64
+Lux.bf16
 ```
 
 ## Element Type Matching

src/Lux.jl

@@ -114,7 +114,7 @@ export BinaryCrossEntropyLoss, BinaryFocalLoss, CrossEntropyLoss, DiceCoeffLoss,
        PoissonLoss, SiameseContrastiveLoss, SquaredHingeLoss
 export GenericLossFunction
 
-export f16, f32, f64
+export f16, f32, f64, bf16
 export match_eltype
 
 export FromFluxAdaptor, FluxLayer

src/helpers/eltype_conversion.jl

@@ -21,14 +21,46 @@ for (fname, ftype) in zip((:f16, :f32, :f64), (Float16, Float32, Float64))
             $($fname)(m)
 
         Converts the `eltype` of `m` *floating point* values to `$($ftype)`.
-        Recurses into structs marked with `Functors.@functor`.
+        To avoid recursion into structs mark them with `Functors.@leaf`.
         """
-        $(fname)(m) = (LuxEltypeAdaptor{$ftype}())(m)
+        $(fname)(m) = (LuxEltypeAdaptor{$(ftype)}())(m)
     end
 end
 
+@static if isdefined(Core, :BFloat16)
+    bf16_docs = """
+        bf16(m)
+
+    Converts the `eltype` of `m` *floating point* values to `BFloat16`.
+    To avoid recursion into structs mark them with `Functors.@leaf`.
+
+    !!! warning
+
+        `BFloat16s.jl` needs to be loaded before using this function.
+
+    !!! tip "Support for `BFloat16`"
+
+        Most Lux operations aren't optimized for `BFloat16` yet. Instead this is meant to be
+        used together with `Reactant.@compile`.
+    """
+
+    bf16(m) = (LuxEltypeAdaptor{Core.BFloat16}())(m)
+else
+    bf16_docs = """
+        bf16(m)
+
+    !!! danger "Not Supported"
+
+        Current Julia version does not support `BFloat16`. Use julia 1.11 or newer.
+    """
+
+    bf16(_) = error("`bf16` is not supported on Julia versions 1.11+")
+end
+
+@doc (bf16_docs) bf16
+
 # Common incorrect usage
-for f in (f16, f32, f64)
+for f in (f16, f32, f64, bf16)
     warn_msg = "$(f) is not meant to be broadcasted like `$(f).(x)` or `x .|> $(f)`, \
                 and this might give unexpected results and could lead to crashes. Directly \
                 use `$(f)` as `$(f)(x)` or `x |> $(f)` instead."