mimiquate/candlex

An Nx backend for candle machine learning framework

candlex's Issues

`Nx.fft`: Implement

Depends on #75

---

https://hexdocs.pm/nx/Nx.html#fft/2

Example failure:

[error] Task #PID<0.516.0> started from SpeechToText.Serving terminating
** (RuntimeError) unsupported Candlex.Backend.fft function
    (candlex 0.1.5) lib/candlex/backend.ex:926: Candlex.Backend.fft/3

`Nx.slice`: Support strides option

`Nx.dot`: support contracting axes in any position for M x N

`Nx.argsort`: implement

Candle Error: the candle crate has not been built with cuda support

Hi, thanks so much for making this package so quickly!

I'm trying to use candlex in an elixir livebook on a host with an A100 and cuda installed (and working via EXLA backend).

The livebook looks like this

# setup cell
Mix.install(
  [
    {:kino_bumblebee, "~> 0.4.0"},
    {:exla, ">= 0.0.0"},
    {:candlex, "~> 0.1.4"}
  ],
  config: [nx: [default_backend: Candlex.Backend]]
)

# subsequent cell

repo =
  {:hf, "codellama/CodeLlama-7b-hf"}

{:ok, model_info} = Bumblebee.load_model(repo, backend: {Candlex.Backend, device: :cuda})
{:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
{:ok, generation_config} = Bumblebee.load_generation_config(repo)

generation_config = Bumblebee.configure(generation_config, max_new_tokens: 100)

serving =
  Bumblebee.Text.generation(model_info, tokenizer, generation_config,
    compile: [batch_size: 1, sequence_length: 1028],
    stream: false,
    defn_options: [compiler: EXLA, lazy_transfers: :never],
    preallocate_params: true
  )

Kino.start_child({Nx.Serving, name: Codellama, serving: serving})

prompt = "<PRE> #fibonnaci in elixir <SUF>Fibby.fib(5) <MID>\n"

Nx.Serving.batched_run(CodeLlama, prompt)

I get the error Candle Error: the candle crate has not been built with cuda support when calling Bumblebee.load_model. How would I go about installing the correct crate? Sorry, I'm a total rust and rustler newbie.

Any help is much appreciated <3

Other than that, would you expect the above to work and with performance on-par with the EXLA backend?

`Nx.put_slice`: less restrictive with shapes

Example of currently unsupported shapes:

        ** (RuntimeError) Unsupported put_slice shapes, tensor={1, 3001, 400} and slice={1, 1, 400}. All-but-last dimensions in slice need to be equal to corresponding dimension in tensor.
            (candlex 0.1.5) lib/candlex/backend.ex:423: Candlex.Backend.put_slice/4

See commented test cases in

candlex/test/candlex_test.exs

Lines 1855 to 1923 in fd21e30

	test "put_slice" do
	t([0, 1, 2, 3, 4])
	|> Nx.put_slice([2], Nx.tensor([5, 6]))
	|> assert_equal(t([0, 1, 5, 6, 4]))
	t([[1, 2, 3], [4, 5, 6]])
	|> Nx.put_slice([0, 0], t([[7, 8, 9], [10, 11, 12]]))
	|> assert_equal(
	t([
	[7, 8, 9],
	[10, 11, 12]
	])
	)
	t([[1, 2, 3], [4, 5, 6]])
	|> Nx.put_slice([0, 1], t([[7, 8], [9, 10]]))
	|> assert_equal(
	t([
	[1, 7, 8],
	[4, 9, 10]
	])
	)
	# t([[1, 2, 3], [4, 5, 6]])
	# |> Nx.put_slice([t(0), t(1)], t([[10.0, 11.0]]))
	# |> assert_equal(t(
	# [
	# [1.0, 10.0, 11.0],
	# [4.0, 5.0, 6.0]
	# ]
	# ))
	# t([[1, 2, 3], [4, 5, 6]])
	# |> Nx.put_slice([1, 1], t([[7, 8], [9, 10]]))
	# |> assert_equal(t(
	# [
	# [1, 7, 8],
	# [4, 9, 10]
	# ]
	# ))
	t([
	[
	[1, 2],
	[3, 4]
	],
	[
	[4, 5],
	[6, 7]
	]
	])
	|> Nx.put_slice([0, 0, 1], t([[[8], [9]], [[10], [11]]]))
	|> assert_equal(
	t([
	[
	[1, 8],
	[3, 9]
	],
	[
	[4, 10],
	[6, 11]
	]
	])
	)
	# t([[[1, 2], [3, 4]]])
	# |> Nx.put_slice([0, 0, 0], t([[[10, 11]]]))
	# |> assert_equal(t([[[10, 11], [3, 4]]]))
	end

`Nx.take`: support indices of rank > 1

candlex/test/candlex_test.exs

Lines 2203 to 2218 in 9e133fe

	# t([[1, 2], [11, 12]])
	# |> Nx.take(t([[0, 0], [1, 1], [0, 0]]), axis: 1)
	# |> assert_equal(t(
	# [
	# [
	# [1, 1],
	# [2, 2],
	# [1, 1]
	# ],
	# [
	# [11, 11],
	# [12, 12],
	# [11, 11]
	# ]
	# ]
	# ))

`Nx.dot`: support batched axes

Precompiled NIF binaries targeting CUDA

Follow up on #9

`Nx.LinAlg`: implement related callbacks

• cholesky/2
• eigh/3
• lu/3
• qr/3
• triangular_solve/4

`Nx.dot`: complete implementation

Nx.dot/2:

Returns the dot product of two tensors.

Given a and b, computes the dot product according to the following rules:

• If both a and b are scalars, it is equivalent to a * b.
• If a is a scalar and b is a tensor, it is equivalent to Nx.multiply(a, b).
• If a is a tensor and b is a scalar, it is equivalent to Nx.multiply(a, b).
• If both a and b are 1-D tensors (vectors), it is the sum of the element-wise product between a and b. The lengths of a and b must be equal.
• If both a and b are 2-D tensors (matrices), it is equivalent to matrix-multiplication.
• If either a or b is a 1-D tensor, and the other is an n-D tensor, it is the sum of the element-wise product along the last axis of a or b. The length of the 1-D tensor must match the last dimension of the n-D tensor.
• If a is an n-D tensor and b is an m-D tensor, it is the sum of the element-wise product along the last axis of a and the second-to-last axis of b. The last dimension of a must match the second-to-last dimension of b.

Cases:

• If both a and b are scalars, it is equivalent to a * b.
• If a is a scalar and b is a tensor, it is equivalent to Nx.multiply(a, b).
• If a is a tensor and b is a scalar, it is equivalent to Nx.multiply(a, b).
• If both a and b are 1-D tensors (vectors), it is the sum of the element-wise product between a and b. The lengths of a and b must be equal.
• If both a and b are 2-D tensors (matrices), it is equivalent to matrix-multiplication.
• If either a or b is a 1-D tensor, and the other is an n-D tensor, it is the sum of the element-wise product along the last axis of a or b. The length of the 1-D tensor must match the last dimension of the n-D tensor. (#17)
• If a is an n-D tensor and b is an m-D tensor, it is the sum of the element-wise product along the last axis of a and the second-to-last axis of b. The last dimension of a must match the second-to-last dimension of b. (#51)

---

Right now is partially implemented given the limitation of having only matmul in candle-core.

See partial test coverage we have in

candlex/test/candlex_test.exs

Lines 500 to 655 in b55a50b

	test "dot/2" do
	# Dot product of scalars
	Nx.dot(5, 5)
	|> assert_equal(t(25))
	Nx.dot(-2.0, 5.0)
	|> assert_equal(t(-10.0))
	Nx.dot(2, 2.0)
	|> assert_equal(t(4.0))
	# Dot product of vectors
	t([1, 2, 3])
	|> Nx.dot(t([4, 5, 6]))
	|> assert_equal(t(32))
	t([1.0, 2, 3])
	|> Nx.dot(t([1, 2, 3]))
	|> assert_equal(t(14.0))
	# Dot product of matrices (2-D tensors)
	# TODO: Candle matmul doesn't support integers yet
	# t([[1, 2, 3], [4, 5, 6]])
	# |> Nx.dot(t([[7, 8], [9, 10], [11, 12]]))
	# |> assert_equal(t(
	# [
	# [58, 64],
	# [139, 154]
	# ]
	# ))
	t([[1.0, 2, 3], [4, 5, 6]])
	|> Nx.dot(t([[7, 8], [9, 10], [11, 12]]))
	|> assert_equal(
	t([
	[58.0, 64],
	[139, 154]
	])
	)
	# Dot product of vector and n-D tensor
	t([[0.0]])
	|> Nx.dot(t([55.0]))
	|> assert_equal(t([0.0]))
	t([[[1.0, 2], [3, 4]], [[5, 6], [7, 8]]])
	|> Nx.dot(t([5, 10]))
	|> assert_equal(
	t([
	[25.0, 55],
	[85, 115]
	])
	)
	# t([5.0, 10], names: [:x])
	# |> Nx.dot(t([[1.0, 2, 3], [4, 5, 6]], names: [:i, :j]))
	# |> assert_equal(t(
	# [45, 60, 75]
	# ))
	# t([[[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]], names: [:shard, :batch, :x, :y, :z])
	# |> Nx.dot(t([2.0, 2.0], names: [:data]))
	# |> assert_equal(t(
	# [
	# [
	# [
	# [6.0, 14.0],
	# [22.0, 30.0]
	# ]
	# ]
	# ]
	# ))
	# Dot product of n-D and m-D tensors
	# t([[[1.0, 2, 3], [4, 5, 6], [7, 8, 9]], [[1, 2, 3], [4, 5, 6], [7, 8, 9]]], names: [:x, :y, :z])
	# |> Nx.dot(t([[[1.0, 2, 3], [3, 4, 5], [5, 6, 7]]], names: [:i, :j, :k]))
	# |> assert_equal(t(
	# [
	# [
	# [
	# [22, 28, 34]
	# ],
	# [
	# [49, 64, 79]
	# ],
	# [
	# [76, 100, 124]
	# ]
	# ],
	# [
	# [
	# [22, 28, 34]
	# ],
	# [
	# [49, 64, 79]
	# ],
	# [
	# [76, 100, 124]
	# ]
	# ]
	# ]
	# ))
	end
	test "dot/6" do
	# Contracting along axes
	t1 = t([[1.0, 2], [3, 4]], names: [:x, :y])
	t2 = t([[10.0, 20], [30, 40]], names: [:height, :width])
	t1
	|> Nx.dot([0], [], t2, [0], [])
	|> assert_equal(
	t([
	[100, 140],
	[140, 200]
	])
	)
	# TODO:
	t1
	|> Nx.dot([0], [], t2, [1], [])
	|> assert_equal(
	t([
	[70, 150],
	[100, 220]
	])
	)
	t1
	|> Nx.dot([1], [], t2, [0], [])
	|> assert_equal(
	t([
	[70, 100],
	[150, 220]
	])
	)
	# t1
	# |> Nx.dot([1], [], t2, [1], [])
	# |> assert_equal(t(
	# [
	# [50, 110],
	# [110, 250]
	# ]
	# ))
	# t1
	# |> Nx.dot([0, 1], [], t2, [0, 1], [])
	# |> assert_equal(t(300))
	end

.

Better support for f16 and bf16

`Nx.window_max`: support `:same` padding

Consider rustler_precompiled to avoid dependency on the rust toolchain

`Nx.window_*`: more complete support for window operations

Support complex numbers (c64)

https://hexdocs.pm/nx/Nx.html#complex/2

`Nx.pad`: support negative padding

`Nx.pad`: support tensors of rank > 1

`Nx.argsort`: support tensors of rank = 2

`Nx.window_sum`: implement