Increase performance with tf.function decorator

[julesdesir]

Description of the goal of the PR

Description: increase performance with tf.function decorator

Changes this PR introduces (fill it before implementation)

• : ENH: make @tf.function work with train_step and compute_batch_loss
• : MAINT: small fix in notebook
• : TST: small fix in test_get_negative_samples

Checklist before requesting a review

• I have commented my code, particularly in hard-to-understand areas
• I have typed my code
• I have created / updated the docstrings
• I have updated the README, if relevant
• I have updated the requirements files if a new package is used
• I have tested my code
• The CI pipeline passes
• I have performed a self-review of my code

[gemini-code-assist]

Important

The terms of service for this installation has not been accepted. Please ask the Organization owners to visit the Gemini Code Assist Admin Console to sign it.

[github-actions]

Coverage Report for Python 3.9

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	48	0	100%
choice_learn/basket_models
__init__.py	3	0	100%
dataset.py	117	4	97%	71–74
preprocessing.py	94	78	17%	43–45, 128–364
shopper.py	358	27	92%	165, 194, 343, 363, 378, 381, 395, 684–688, 781–785, 883–887, 1218, 1297, 1335–1336, 1440–1441, 1517–1518
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	647	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	239	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 577
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	393	5	99%	43–44, 154–155, 715
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	254	12	95%	144, 186, 283, 302, 342, 349, 378, 397, 428–429, 438–439
baseline_models.py	49	0	100%
conditional_logit.py	236	21	91%	46, 49, 51, 82, 85, 88–92, 95–99, 133, 298, 335, 392, 467–473, 598, 632, 739, 743
halo_mnl.py	124	2	98%	186, 374
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	236	236	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
TOTAL	4791	700	85%

Tests	Skipped	Failures	Errors	Time
165	0 💤	0 ❌	0 🔥	4m 48s ⏱️

[github-actions]

Coverage Report for Python 3.11

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	48	0	100%
choice_learn/basket_models
__init__.py	3	0	100%
dataset.py	117	4	97%	71–74
preprocessing.py	94	78	17%	43–45, 128–364
shopper.py	358	27	92%	165, 194, 343, 363, 378, 381, 395, 684–688, 781–785, 883–887, 1218, 1297, 1335–1336, 1440–1441, 1517–1518
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	647	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	239	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 577
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	393	4	99%	39, 154–155, 715
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	254	12	95%	144, 186, 283, 302, 342, 349, 378, 397, 428–429, 438–439
baseline_models.py	49	0	100%
conditional_logit.py	236	21	91%	46, 49, 51, 82, 85, 88–92, 95–99, 133, 298, 335, 392, 467–473, 598, 632, 739, 743
halo_mnl.py	124	18	85%	186, 341, 360, 364–380
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
TOTAL	4793	717	85%

Tests	Skipped	Failures	Errors	Time
165	0 💤	1 ❌	0 🔥	5m 12s ⏱️

[github-actions]

Coverage Report for Python 3.10

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	48	0	100%
choice_learn/basket_models
__init__.py	3	0	100%
dataset.py	117	4	97%	71–74
preprocessing.py	94	78	17%	43–45, 128–364
shopper.py	358	27	92%	165, 194, 343, 363, 378, 381, 395, 684–688, 781–785, 883–887, 1218, 1297, 1335–1336, 1440–1441, 1517–1518
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	647	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	239	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 577
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	393	4	99%	39, 154–155, 715
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	254	12	95%	144, 186, 283, 302, 342, 349, 378, 397, 428–429, 438–439
baseline_models.py	49	0	100%
conditional_logit.py	236	21	91%	46, 49, 51, 82, 85, 88–92, 95–99, 133, 298, 335, 392, 467–473, 598, 632, 739, 743
halo_mnl.py	124	18	85%	186, 341, 360, 364–380
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
TOTAL	4793	717	85%

Tests	Skipped	Failures	Errors	Time
165	0 💤	1 ❌	0 🔥	5m 12s ⏱️

[github-actions]

Coverage Report for Python 3.12

File	Stmts	Miss	Cover	Missing
choice_learn
__init__.py	2	0	100%
tf_ops.py	48	0	100%
choice_learn/basket_models
__init__.py	3	0	100%
dataset.py	117	4	97%	71–74
preprocessing.py	94	78	17%	43–45, 128–364
shopper.py	358	27	92%	165, 194, 343, 363, 378, 381, 395, 684–688, 781–785, 883–887, 1218, 1297, 1335–1336, 1440–1441, 1517–1518
choice_learn/basket_models/utils
__init__.py	0	0	100%
permutation.py	22	1	95%	37
choice_learn/data
__init__.py	3	0	100%
choice_dataset.py	647	33	95%	198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
indexer.py	239	23	90%	20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 577
storage.py	161	6	96%	22, 33, 51, 56, 61, 71
store.py	72	72	0%	3–275
choice_learn/datasets
__init__.py	4	0	100%
base.py	393	4	99%	39, 154–155, 715
expedia.py	102	83	19%	37–301
tafeng.py	49	0	100%
choice_learn/datasets/data
__init__.py	0	0	100%
choice_learn/models
__init__.py	14	2	86%	15–16
base_model.py	254	12	95%	144, 186, 283, 302, 342, 349, 378, 397, 428–429, 438–439
baseline_models.py	49	0	100%
conditional_logit.py	236	21	91%	46, 49, 51, 82, 85, 88–92, 95–99, 133, 298, 335, 392, 467–473, 598, 632, 739, 743
halo_mnl.py	124	2	98%	186, 374
latent_class_base_model.py	286	39	86%	55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
latent_class_mnl.py	62	6	90%	257–261, 296
learning_mnl.py	67	3	96%	157, 182, 188
nested_logit.py	291	12	96%	55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
reslogit.py	132	6	95%	285, 360, 369, 374, 382, 432
rumnet.py	236	3	99%	748–751, 982
simple_mnl.py	139	6	96%	167, 275, 347, 355, 357, 359
tastenet.py	94	3	97%	142, 180, 188
choice_learn/toolbox
__init__.py	0	0	100%
assortment_optimizer.py	27	6	78%	28–30, 93–95, 160–162
gurobi_opt.py	238	238	0%	3–675
or_tools_opt.py	230	11	95%	103, 107, 296–305, 315, 319, 607, 611
TOTAL	4793	701	85%

Tests	Skipped	Failures	Errors	Time
165	0 💤	0 ❌	0 🔥	6m 7s ⏱️