Nx.put_slice: less restrictive with shapes

[grzuy]

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

[grzuy]

I think something like this could help huggingface/candle#1351.