change import from tensorflow.keras to keras. #562

Job	Run time
test (ubuntu-latest, 3.10, jax)	24s
test (ubuntu-latest, 3.10, numpy)	22s
test (ubuntu-latest, 3.10, tensorflow)	46s
test (ubuntu-latest, 3.10, torch)	25s
test (ubuntu-latest, 3.11, jax)	27s
test (ubuntu-latest, 3.11, numpy)	27s
test (ubuntu-latest, 3.11, tensorflow)	24s
test (ubuntu-latest, 3.11, torch)	26s
test (windows-latest, 3.10, jax)	26s
test (windows-latest, 3.10, numpy)	23s
test (windows-latest, 3.10, tensorflow)	25s
test (windows-latest, 3.10, torch)	32s
test (windows-latest, 3.11, jax)	22s
test (windows-latest, 3.11, numpy)	23s
test (windows-latest, 3.11, tensorflow)	27s
test (windows-latest, 3.11, torch)	22s
	7m 1s

Provide feedback