Improve free symbols search for large tensors

[kinianlo]

Problem

tensor.Box.free_symbols very slow for large tensors. This is because a isinstance is called for each entry in the tensor.

Solution

This PR optimise the free_symbols method such that numerical arrays are not recursively searched for symbols.

import time

from lambeq import Symbol
from lambeq.backend.tensor import Box, Dim
import numpy as np

sym = Symbol('x')

d = 1000
box_with_sym = Box('box_sym', Dim(d), Dim(d))
box_without_sym = Box('box_no_sym', Dim(d), Dim(d))

box_with_sym.data = np.random.rand(d, d).astype(object)
box_without_sym.data = np.random.rand(d, d)
box_with_sym.data[0, 0] = sym

start = time.time()
symbols = box_with_sym.free_symbols
end = time.time()
print(f"Time taken: {end - start:.6f} seconds")

start = time.time()
symbols = box_without_sym.free_symbols
end = time.time()
print(f"Time taken: {end - start:.6f} seconds")

without this PR

(With symbol) Time taken: 0.563968 seconds
(Without symbol) Time taken: 0.555081 seconds

with this PR

(With symbol) Time taken: 0.559249 seconds
(Without symbol) Time taken: 0.000064 seconds

[dimkart]

@kinianlo Thanks for the PR, looks ready to be merged as soon as CI passes.

[kinianlo]

@kinianlo Thanks for the PR, looks ready to be merged as soon as CI passes.

Thanks @dimkart. I was trying to look into the CI errors but they don't seem to be related to the file changed in this PR.

[neiljdo]

LGTM! Thank you @kinianlo for this.